simple TOF model test actually works - o, glorious day..

tests/mcModelIntegration.py

@@ -5,6 +5,8 @@
 #
 # intended to demo and develop MC-based integration of models
 # needed to numerically evaluate likelihoods / probabilities
+# 
+# all the 'magic' here is in the functions lnlike and generateModelData
 
 import numpy as np
 from numpy import inf
@@ -32,6 +34,13 @@
 distance_cellLength = 2.86 # cm, length of gas cell
 distance_zeroDegLength = 3.81 # cm, length of 0deg detector
 
+##############
+# vars for binning of TOF 
+tof_nBins = 25
+tof_minRange = 175.0
+tof_maxRange = 200.0
+tof_range = (tof_minRange,tof_maxRange)
+
 
 def getDDneutronEnergy(deuteronEnergy, labAngle = 0):
     """Get the energy of neutrons produced by DDN reaction
@@ -100,7 +109,7 @@ def lnlike(params, observables, nDraws=2000000):
     """
     #print('checking type ({}) and length ({}) of params in lnlikefxn'.format(type(params),len(params)))
     evalData=generateModelData(params, nDraws)
-    evalHist, evalBinEdges = np.histogram(evalData[:,3], 25, (175.0,200.0),
+    evalHist, evalBinEdges = np.histogram(evalData[:,3], tof_nBins, tof_range,
                                           density=True)
     logEvalHist = np.log(evalHist)
     #print(logEvalHist)
@@ -179,7 +188,8 @@ def lnlike(params, observables, nDraws=2000000):
 plot.xlabel('Location in cell (cm)')
 plot.show()
 
-histTOFdata, tof_bin_edges = np.histogram(fakeDataSet[:,3], 25, (175.0,200.0),
+histTOFdata, tof_bin_edges = np.histogram(fakeDataSet[:,3], tof_nBins, 
+                                          tof_range,
                                           density=True)
 plot.figure(20)
 plot.scatter(tof_bin_edges[:-1], histTOFdata, color='k')
@@ -201,8 +211,8 @@ def lnlike(params, observables, nDraws=2000000):
 
 
 # make our vector of 'n'
-observedVectorN, observed_bin_edges = np.histogram(fakeObsData[:,3], 25,
-                                                   (175.0,200.0))
+observedVectorN, observed_bin_edges = np.histogram(fakeObsData[:,3], tof_nBins,
+                                                   tof_range)
 
 
 loghist = np.log(histTOFdata)
@@ -255,7 +265,7 @@ def lnlike(params, observables, nDraws=2000000):
 
 plot.figure()
 plot.subplot(221)
-plot.hist(fakeObsData[:,3], 25, (175.0,200.0))
+plot.hist(fakeObsData[:,3], tof_nBins, tof_range)
 plot.subplot(223)
 plot.scatter(multiplier,nll_sigmas)
 plot.xlabel('fraction sigma TRUE')

tests/simpleTOFmodel.py

@@ -7,6 +7,7 @@
 # doesn't include all the realistic features needed to do actual analysis
 
 import numpy as np
+from numpy import inf
 from scipy.integrate import quad
 import scipy.optimize as optimize
 import matplotlib.pyplot as plot
@@ -31,6 +32,13 @@
 distance_cellLength = 2.86 # cm, length of gas cell
 distance_zeroDegLength = 3.81 # cm, length of 0deg detector
 
+##############
+# vars for binning of TOF 
+tof_nBins = 25
+tof_minRange = 175.0
+tof_maxRange = 200.0
+tof_range = (tof_minRange,tof_maxRange)
+
 
 def getDDneutronEnergy(deuteronEnergy, labAngle = 0):
     """Get the energy of neutrons produced by DDN reaction
@@ -58,7 +66,69 @@ def getTOF(mass, energy, distance):
     return tof
 
 
-def evaluateModel()
+def generateModelData(params, nSamples):
+    """Generate some fake data from our mode with num samples = nSamples
+    params is an array of the parameters, [e0, e1, sigma]
+    Returns a tuple of len(nSamples), [x, ed, en, tof]
+    """
+    initialEnergy, eLoss, sigma = params
+    data_x=np.random.uniform(low=0.0, high=distance_cellLength, size=nSamples)
+    data_ed= np.random.normal(loc=initialEnergy + eLoss*data_x, 
+                              scale=sigma)
+    data_en = getDDneutronEnergy(data_ed)
+
+    neutronDistance = distance_cellToZero + (distance_cellLength - data_x)
+    neutronTOF = getTOF(mass_neutron, data_en, neutronDistance)
+    effectiveDenergy = (initialEnergy + data_ed)/2
+    deuteronTOF = getTOF( mass_deuteron, effectiveDenergy, data_x )
+    data_tof = neutronTOF + deuteronTOF
+
+    data = np.column_stack((data_x,data_ed,data_en,data_tof))
+    return data
+
+def lnlike(params, observables, nDraws=1000000):
+    """Evaluate the log likelihood given a set of params and observables
+    Observables is a vector; a histogrammed time distribution
+    Params is a list of [initial D energy, E_D loss, sigma]
+    nDraws is the number of points drawn from (energy,location) distribution\
+    which are used to produce the PDF evaluations at different TOFs
+    """
+    #print('checking type ({}) and length ({}) of params in lnlikefxn'.format(type(params),len(params)))
+    evalData=generateModelData(params, nDraws)
+    evalHist, evalBinEdges = np.histogram(evalData[:,3], tof_nBins, tof_range,
+                                          density=True)
+    logEvalHist = np.log(evalHist)
+    #print(logEvalHist)
+    # find what TOFs have zero observed data
+    # we'll use this to handle cases where we might wind up with -inf*0
+    # likelihood is fine if PDF is 0 somewhere where no data is found
+    # without checks though, ln(PDF=0)=-inf, -inf*0 = nan
+    # however, if PDF is 0 (lnPDF=-inf) where there IS data, lnL should be -inf
+    zeroObservedIndices = np.where(observables == 0)[0]
+    for idx in zeroObservedIndices:
+        if logEvalHist[idx] == -inf:
+            logEvalHist[zeroObservedIndices] = 0
+
+    loglike = np.dot(logEvalHist,observables)
+    return loglike
+
+
+
+def lnprior(theta):
+    e_0, e_1, sigma = theta
+    if 800 < e_0 < 1200 and -200 <e_1<0 and 10 < sigma < 100:
+        return 0
+    return -inf
+
+def lnprob(theta, observables):
+    """Evaluate the log probability
+    theta is a list of the model parameters
+    observables is, in this case, a histogram of TOF values
+    """
+    prior = lnprior(theta)
+    if not np.isfinite(prior):
+        return -inf
+    return prior + lnlike(theta, observables)
 
 # mp_* are model parameters
 # *_t are 'true' values that go into our fake data
@@ -69,39 +139,28 @@ def evaluateModel()
 
 # generate fake data
 nSamples = 10000
-data_model_x = np.random.uniform(low=0.0, high=distance_cellLength, 
-                                 size=nSamples)
-data_model_ed = np.random.normal(loc=mp_initialEnergy_t + 
-                                 mp_loss0_t * data_model_x,
-                                 scale=mp_sigma_t)
-data_model_en = getDDneutronEnergy(data_model_ed, 0.0)
-
-# we've got now a neutron energy distribution over the length of the cell
-# let's make some TOF data from that
-neutronDistance = distance_cellToZero + (distance_cellLength - data_model_x)
-neutronTOF = getTOF(mass_neutron, data_model_en, neutronDistance)
-effectiveDenergy = (mp_initialEnergy_t + data_model_ed)/2
-deuteronTOF = getTOF( mass_deuteron, effectiveDenergy, data_model_x )
-modelTOFdata = neutronTOF + deuteronTOF
+fakeData = generateModelData([mp_initialEnergy_t,mp_loss0_t,mp_sigma_t],
+                             nSamples)
+
 
 
 # plot the fake data...
 plot.figure(1)
-plot.scatter(data_model_x, data_model_en, color='k', alpha=0.3)
+plot.scatter(fakeData[:,0], fakeData[:,2], color='k', alpha=0.3)
 plot.xlabel('Cell location (cm)')
 plot.ylabel('Neutron energy (keV)')
 plot.show()
 
 
 # plot the TOF 
 plot.figure(2)
-plot.hist(modelTOFdata, bins=50)
+plot.hist(fakeData[:,3], bins=50)
 plot.xlabel('TOF (ns)')
 plot.show()
 
 # plot the TOF vs x location
 plot.figure(3)
-plot.scatter(data_model_en,modelTOFdata, color='k', alpha=0.3)
+plot.scatter(fakeData[:,2],fakeData[:,3], color='k', alpha=0.3)
 plot.xlabel('Neutron energy (keV)' )
 plot.ylabel('TOF (ns)')
 plot.show()
@@ -114,4 +173,53 @@ def evaluateModel()
 # unfortunately, even a regular old PDF is a hideously complicated thing
 # no real chance of an analytical approach
 # but we can NUMERICALLY attempt to do things
-# for a given 
+
+nll = lambda *args: -lnlike(*args)
+
+
+observedTOF, observed_bin_edges = np.histogram(fakeData[:,3], 
+                                               tof_nBins, tof_range)
+
+minimizedNLL = optimize.minimize(nll, [1080,
+                                       mp_loss0_t *1.2,mp_sigma_t *1.05], 
+                                       args=observedTOF, method='Nelder-Mead',
+                                       tol=1.0)
+
+print(minimizedNLL)
+
+
+nDim, nWalkers = 3, 100
+
+e0, e1, sigma = minimizedNLL["x"]
+
+p0 = [[e0,e1,sigma] + 1e-2 * np.random.randn(nDim) for i in range(nWalkers)]
+sampler = emcee.EnsembleSampler(nWalkers, nDim, lnprob, 
+                                kwargs={'observables': observedTOF}, 
+                                threads=8)
+
+#sampler.run_mcmc(p0, 500)
+# run with progress updates..
+for i, samplerResult in enumerate(sampler.sample(p0, iterations=500)):
+    if (i+1)%10 == 0:
+        print("{0:5.1%}".format(float(i)/500))
+
+plot.figure()
+plot.subplot(311)
+plot.plot(sampler.chain[:,:,0].T,'-',color='k',alpha=0.2)
+plot.ylabel('Initial energy (keV)')
+plot.subplot(312)
+plot.plot(sampler.chain[:,:,1].T,'-',color='k',alpha=0.2)
+plot.ylabel('Energy loss (keV/cm)')
+plot.subplot(313)
+plot.plot(sampler.chain[:,:,2].T,'-',color='k',alpha=0.2)
+plot.ylabel('Sigma (keV)')
+plot.xlabel('Step')
+plot.show()
+
+
+samples = sampler.chain[:,200:,:].reshape((-1,nDim))
+import corner as corn
+cornerFig = corn.corner(samples,labels=["$E_0$","$E_1$","$\sigma$"],
+                        truths=[mp_initialEnergy_t, mp_loss0_t, mp_sigma_t],
+                        quantiles=[0.16,0.5,0.84], show_titles=True,
+                        title_kwargs={'fontsize': 12})