diff options
| -rw-r--r-- | likelihood.py | 74 |
1 files changed, 74 insertions, 0 deletions
diff --git a/likelihood.py b/likelihood.py new file mode 100644 index 0000000..f0323fe --- /dev/null +++ b/likelihood.py @@ -0,0 +1,74 @@ +import numpy as np +from histogram import HistogramDD +from uncertainties import ufloat, umath +from itertools import izip, compress + +class Likelihood(object): + "Class to evaluate likelihoods for detector events." + def __init__(self, sim, event=None): + """ + Args: + - sim: chroma.sim.Simulation + The simulation object used to simulate events and build pdfs. + - event: chroma.event.Event, *optional* + The detector event being reconstructed. If None, you must call + set_event() before eval(). + """ + self.sim = sim + + if event is not None: + self.set_event(event) + + self.pdfs = dict(zip(sim.detector.pmtids,[HistogramDD((100,10), range=[(0.0,1000.0),(-0.5,9.5)])]*len(sim.detector.pmtids))) + + def set_event(self, event): + "Set the detector event being reconstructed." + self.event = event + + def eval(self, vertex_generator, nevals): + """ + Return the negative log likelihood that the detector event set in the + constructor or by set_event() was the result of a particle generated + by `vertex_generator`. + """ + for pdf in self.pdfs.values(): + pdf.reset() + + for ev in self.sim.simulate(nevals, vertex_generator): + for i, t, q in compress(izip(range(ev.hits.hit.size),ev.hits.t,ev.hits.q),ev.hits.hit): + self.pdfs[i].fill((ev.hits.t[i],ev.hits.q[i])) + + for pdf in self.pdfs.values(): + if pdf.nentries > 0: + pdf.normalize() + + log_likelihood = ufloat((0,0)) + + for i, t, q in compress(izip(range(self.event.channels.hit.size),self.event.channels.t,self.event.channels.q),self.event.channels.hit): + probability = self.pdfs[i].ueval((t,q)).item() + + if probability.nominal_value == 0.0: + if self.pdfs[i].nentries > 0: + probability = ufloat([0.5/self.pdfs[i].nentries]*2) + else: + probability = ufloat([1.0/len(self.pdfs[i].bincenters)]*2) + + log_likelihood += umath.log(probability) + + return -log_likelihood + +if __name__ == '__main__': + from chroma import detectors + from chroma.sim import Simulation + from chroma.optics import water + from chroma.generator import constant_particle_gun + + detector = detectors.find('minilbne') + sim = Simulation(detector, water) + + event = sim.simulate(1, constant_particle_gun('e-',(0,0,0),(1,0,0),100.0)).next() + + likelihood = Likelihood(sim, event) + + for x in np.linspace(-1.0, 1.0, 20): + print 'x = %f, %s' % (x, likelihood.eval(constant_particle_gun('e-',(x,0,0),(1,0,0),100.0),100)) |