Merge branch 'examples'

lorenzennio · Nov 23, 2023 · c6c6793 · c6c6793
2 parents 2c5bd0e + 7957c5d
commit c6c6793
Show file tree

Hide file tree

Showing 8 changed files with 3,890 additions and 0 deletions.
diff --git a/examples/knunu.ipynb b/examples/knunu.ipynb
diff --git a/examples/knunu_sample.ipynb b/examples/knunu_sample.ipynb
diff --git a/examples/knunu_sample.py b/examples/knunu_sample.py
@@ -0,0 +1,34 @@
+import numpy as np
+from publik import modifier
+import knunu_utils
+
+from Bayesian_pyhf import infer
+from Bayesian_pyhf import prepare_inference
+import pymc as pm
+
+null = knunu_utils.null_pred()
+alt = knunu_utils.alt_pred()
+
+model, alt_yields = modifier.load('knunu_model.json', alt.distribution, null.distribution, return_data=True)
+
+# Perform the sampling
+unconstr_priors = {
+    'mu':  {'type': 'Normal_Unconstrained', 'mu': [1.], 'sigma': [1e-10]},
+    'cvl': {'type': 'Uniform_Unconstrained', 'lower': [2.], 'upper': [10.]},
+    'csl': {'type': 'Uniform_Unconstrained', 'lower': [0.], 'upper': [5.]},
+    'ctl': {'type': 'Uniform_Unconstrained', 'lower': [0.], 'upper': [5.]}
+}
+
+priorDict_conjugate = prepare_inference.build_priorDict(model, unconstr_priors)
+priorDict_conjugate
+
+n_draws = 10000
+with infer.model(model, unconstr_priors, alt_yields):
+    # step = pm.Metropolis()
+    post_data = pm.sample(draws=n_draws)#, step=step, tune=1000)
+    post_pred = pm.sample_posterior_predictive(post_data)
+    prior_pred = pm.sample_prior_predictive(n_draws)
+
+post_data.to_json( 'samples/nuts_post_data.json')
+post_pred.to_json( 'samples/nuts_post_pred.json')
+prior_pred.to_json('samples/nuts_prior_pred.json')
diff --git a/examples/knunu_utils.py b/examples/knunu_utils.py
@@ -0,0 +1,140 @@
+import numpy as np
+import eos
+import numbers
+
+def analysis():
+    """
+    Specify the likelihoods and FF parameter ranges 
+    
+    Returns:
+        EOS analysis instance
+    """
+
+    # form factor expansion f+_0,1,2 are expansion parameters up to 2nd order
+    # there is no f0_0 because of a constriant which removes one parameter
+
+    parameters = [
+        0.33772497529184886, -0.87793473613271, -0.07935870922121949, 
+        0.3719622997220613, 0.07388594710238389, 0.327935912834808, 
+        -0.9490004115927961, -0.23146429907794228
+        ]
+    paramerror = [
+        0.010131234226468245, 0.09815140228051167, 0.26279803480131697, 
+        0.07751034526769873, 0.14588095119443809, 0.019809720318176644, 
+        0.16833757660616938, 0.36912754148836896
+        ]
+    sigma = 15
+    analysis_args = {
+        'priors': [
+            { 'parameter': 'B->K::alpha^f+_0@BSZ2015', 'min': parameters[0]-sigma*paramerror[0], 'max': parameters[0]+sigma*paramerror[0], 'type': 'uniform' },
+            { 'parameter': 'B->K::alpha^f+_1@BSZ2015', 'min': parameters[1]-sigma*paramerror[1], 'max': parameters[1]+sigma*paramerror[1], 'type': 'uniform' },
+            { 'parameter': 'B->K::alpha^f+_2@BSZ2015', 'min': parameters[2]-sigma*paramerror[2], 'max': parameters[2]+sigma*paramerror[2], 'type': 'uniform' },
+            { 'parameter': 'B->K::alpha^f0_1@BSZ2015', 'min': parameters[3]-sigma*paramerror[3], 'max': parameters[3]+sigma*paramerror[3], 'type': 'uniform' },
+            { 'parameter': 'B->K::alpha^f0_2@BSZ2015', 'min': parameters[4]-sigma*paramerror[4], 'max': parameters[4]+sigma*paramerror[4], 'type': 'uniform' },
+            { 'parameter': 'B->K::alpha^fT_0@BSZ2015', 'min': parameters[5]-sigma*paramerror[5], 'max': parameters[5]+sigma*paramerror[5], 'type': 'uniform' },
+            { 'parameter': 'B->K::alpha^fT_1@BSZ2015', 'min': parameters[6]-sigma*paramerror[6], 'max': parameters[6]+sigma*paramerror[6], 'type': 'uniform' },
+            { 'parameter': 'B->K::alpha^fT_2@BSZ2015', 'min': parameters[7]-sigma*paramerror[7], 'max': parameters[7]+sigma*paramerror[7], 'type': 'uniform' }
+        ],
+        'likelihood': [
+            'B->K::f_0+f_++f_T@FLAG:2021A',
+            'B->K::f_0+f_++f_T@HPQCD:2022A'
+        ]
+    }
+
+    analysis = eos.Analysis(**analysis_args)
+    analysis.optimize()
+    return analysis
+
+def efficiency(q2):
+    """
+    Efficiency map adapted from https://journals.aps.org/prl/supplemental/10.1103/PhysRevLett.127.181802/suppl_mat.pdf (Figure 3)
+    Overall scale different, as this can be compensated be number of simulated events.
+    """
+    return np.exp(-0.2*q2)
+
+class null_pred:
+    """
+    Null (SM) prediction
+    """
+    def __init__(self):
+        p = analysis().parameters
+        k = eos.Kinematics({'q2': 0.})
+        o = eos.Options(**{'form-factors': 'BSZ2015', 'model': 'WET'})
+
+        self.kv1 = k['q2']
+
+        self.obs = eos.Observable.make('B->Knunu::dBR/dq2', p, k, o)
+
+    def distribution(self, q2):
+        if isinstance(q2, numbers.Number):
+            self.kv1.set(q2)
+            obs = self.obs.evaluate()
+        else:
+            obs = []
+            for q in q2:
+                self.kv1.set(q)
+                obs.append(self.obs.evaluate())
+
+        return obs
+
+class alt_pred:
+    """
+    Alternative (BSM) prediction
+    """
+    def __init__(self):
+        self.ana = analysis()
+        p = self.ana.parameters
+        k = eos.Kinematics({'q2': 0.})
+        o = eos.Options(**{'form-factors': 'BSZ2015', 'model': 'WET'})
+
+        self.kv1 = k['q2'                      ]
+        self.wc1 = p['sbnunu::Re{cVL}'         ]
+        self.wc2 = p['sbnunu::Re{cSL}'         ]
+        self.wc3 = p['sbnunu::Re{cTL}'         ]
+        self.hv1 = p['B->K::alpha^f+_0@BSZ2015']
+        self.hv2 = p['B->K::alpha^f+_1@BSZ2015']
+        self.hv3 = p['B->K::alpha^f+_2@BSZ2015']
+        self.hv4 = p['B->K::alpha^f0_1@BSZ2015']
+        self.hv5 = p['B->K::alpha^f0_2@BSZ2015']
+        self.hv6 = p['B->K::alpha^fT_0@BSZ2015']
+        self.hv7 = p['B->K::alpha^fT_1@BSZ2015']
+        self.hv8 = p['B->K::alpha^fT_2@BSZ2015']
+
+        self.obs = eos.Observable.make('B->Knunu::dBR/dq2', p, k, o)
+
+    def distribution(self, q2, cvl, csl, ctl, fp0, fp1, fp2, f01, f02, fT0, fT1, fT2):
+        self.wc1.set(cvl)
+        self.wc2.set(csl)
+        self.wc3.set(ctl)
+        self.hv1.set(fp0)
+        self.hv2.set(fp1)
+        self.hv3.set(fp2)
+        self.hv4.set(f01)
+        self.hv5.set(f02)
+        self.hv6.set(fT0)
+        self.hv7.set(fT1)
+        self.hv8.set(fT2)
+
+        if isinstance(q2, numbers.Number):
+            self.kv1.set(q2)
+            obs = self.obs.evaluate()
+        else:
+            obs = []
+            for q in q2:
+                self.kv1.set(q)
+                obs.append(self.obs.evaluate())
+
+        return obs
+
+def parameter_cov(ana):
+    """
+    Get covariance matrix of parameters in EOS analysis object.
+    """
+    pars = []
+    for n in range(0,5):
+        rng = np.random.mtrand.RandomState(74205+n)
+        p, _ = ana.sample(N=5000, stride=5, pre_N=1000, preruns=5, rng=rng)
+        pars += p.tolist()
+    pars = np.array(pars)
+    cov = np.cov(pars.T).tolist()
+    return cov
diff --git a/examples/pilnu.ipynb b/examples/pilnu.ipynb
diff --git a/examples/pilnu_utils.py b/examples/pilnu_utils.py
@@ -0,0 +1,128 @@
+import numpy as np
+import eos
+import numbers
+
+def analysis():
+    """
+    Specify the likelihoods and FF parameter ranges 
+    
+    Returns:
+        EOS analysis instance
+    """
+
+    analysis_args = {
+        'priors': [
+            { 'parameter': 'B->D::alpha^f+_0@BSZ2015', 'min':  0.0, 'max':  1.0, 'type': 'uniform' },
+            { 'parameter': 'B->D::alpha^f+_1@BSZ2015', 'min': -5.0, 'max': +5.0, 'type': 'uniform' },
+            { 'parameter': 'B->D::alpha^f+_2@BSZ2015', 'min': -5.0, 'max': +5.0, 'type': 'uniform' },
+            { 'parameter': 'B->D::alpha^f0_1@BSZ2015', 'min': -5.0, 'max': +5.0, 'type': 'uniform' },
+            { 'parameter': 'B->D::alpha^f0_2@BSZ2015', 'min': -5.0, 'max': +5.0, 'type': 'uniform' }
+        ],
+        'likelihood': [
+            'B->D::f_++f_0@HPQCD:2015A',
+            'B->D::f_++f_0@FNAL+MILC:2015B'
+        ]
+    }
+
+    analysis = eos.Analysis(**analysis_args)
+    analysis.optimize()
+    return analysis
+
+class null_pred:
+    """
+    Null (SM) prediction
+    """
+    def __init__(self):
+        p = eos.Parameters()
+        o = eos.Options({'form-factors': 'BSZ2015', 'l': 'tau', 'model':'WET'})
+        k = eos.Kinematics({'q2': 5.0, 'cos(theta_l)': 0.0,})
+
+        self.kv1 = k['q2']
+        self.kv2 = k['cos(theta_l)']
+
+        self.obs = eos.Observable.make('B->pilnu::d^2BR/dq2/dcos(theta_l)', p, k, o)
+
+
+    def distribution(self, q2, costl):
+        if isinstance(q2, numbers.Number) and isinstance(costl, numbers.Number):
+            self.kv1.set(q2)
+            self.kv2.set(costl)
+            obs = self.obs.evaluate()
+        else:
+            obs = []
+            for q in q2:
+                coslist = []
+                for ct in costl:
+                    self.kv1.set(q)
+                    self.kv2.set(ct)
+                    o = self.obs.evaluate()
+                    coslist.append(o)
+                obs.append(coslist)
+            obs = np.array(obs).T
+
+        return obs
+
+class alt_pred:
+    """
+    Alternative (BSM) prediction
+    """
+    def __init__(self):
+        self.ana = analysis()
+        p = self.ana.parameters
+        o = eos.Options({'form-factors': 'BSZ2015', 'l': 'tau', 'model':'WET'})
+        k = eos.Kinematics({'q2': 5.0, 'cos(theta_l)': 0.0,})
+
+        self.kv1 = k['q2']
+        self.kv2 = k['cos(theta_l)']
+        self.wc1 = p['ubtaunutau::Re{cVL}'     ]
+        self.wc2 = p['ubtaunutau::Re{cSL}'     ]
+        self.wc3 = p['ubtaunutau::Re{cT}'      ]
+        self.hv1 = p['B->D::alpha^f+_0@BSZ2015']
+        self.hv2 = p['B->D::alpha^f+_1@BSZ2015']
+        self.hv3 = p['B->D::alpha^f+_2@BSZ2015']
+        self.hv4 = p['B->D::alpha^f0_1@BSZ2015']
+        self.hv5 = p['B->D::alpha^f0_2@BSZ2015']
+
+        self.obs = eos.Observable.make('B->pilnu::d^2BR/dq2/dcos(theta_l)', p, k, o)
+
+
+    def distribution(self, q2, costl, cvl, csl, ct, fp0, fp1, fp2, f01, f02):
+        self.wc1.set(cvl)
+        self.wc2.set(csl)
+        self.wc3.set(ct )
+        self.hv1.set(fp0)
+        self.hv2.set(fp1)
+        self.hv3.set(fp2)
+        self.hv4.set(f01)
+        self.hv5.set(f02)
+
+        if isinstance(q2, numbers.Number) and isinstance(costl, numbers.Number):
+            self.kv1.set(q2)
+            self.kv2.set(costl)
+            obs = self.obs.evaluate()
+        else:
+            obs = []
+            for q in q2:
+                coslist = []
+                for ct in costl:
+                    self.kv1.set(q)
+                    self.kv2.set(ct)
+                    o = self.obs.evaluate()
+                    coslist.append(o)
+                obs.append(coslist)
+            obs = np.array(obs).T
+
+        return obs
+
+def parameter_cov(ana):
+    """
+    Get covariance matrix of parameters in EOS analysis object.
+    """
+    pars = []
+    for n in range(0,5):
+        rng = np.random.mtrand.RandomState(74205+n)
+        p, _ = ana.sample(N=5000, stride=5, pre_N=1000, preruns=5, rng=rng)
+        pars += p.tolist()
+    pars = np.array(pars)
+    cov = np.cov(pars.T).tolist()
+    return cov