def write_output(self):
# Write the results to a file
from ROOT import TFile
output_file = TFile.Open(self.options().output, "recreate")
output_file.WriteTObject(self._data, self._data.GetName())
gp = self.gen_params()
if gp:
from ROOT import RooArgSet
gpars = RooArgSet()
for p in gp:
gpars.add(p)
gpars = gpars.snapshot(True)
output_file.WriteTObject(gpars, "gen_params")
output_file.Close()
def run(self, **kwargs):
from ROOT import RooArgSet
__check_req_kw__("Observables", kwargs)
__check_req_kw__("Pdf", kwargs)
observables = kwargs.pop("Observables")
obs_set = RooArgSet(*observables)
pdf = kwargs.pop("Pdf")
genPdf = kwargs.pop("GenPdf", pdf)
gen_obs_set = RooArgSet()
for o in list(observables) + list(genPdf.ConditionalObservables()):
gen_obs_set.add(o._target_())
gen_pdf_params = genPdf.getParameters(gen_obs_set).snapshot(True)
genPdf = genPdf.clone(genPdf.GetName() + "_toy_clone")
genPdf.recursiveRedirectServers(gen_pdf_params)
fit_obs_set = RooArgSet()
for o in list(observables) + list(pdf.ConditionalObservables()):
fit_obs_set.add(o._target_())
params = pdf.getParameters(fit_obs_set)
pdf_params = RooArgSet()
for p in params:
if p.isConstant():
continue
pdf_params.add(p)
## for param in pdf_params:
## if param.GetName() not in ['Gamma', 'dGamma']:
## param.setConstant()
self._gen_params = pdf_params.snapshot(True)
# Make another ArgSet to put the fit results in
result_params = RooArgSet(pdf_params, "result_params")
transform = self.transform()
if transform:
trans_params = transform.gen_params(gen_obs_set)
for p in trans_params:
result_params.add(p)
# Some extra numbers of interest
from ROOT import RooRealVar
NLL = RooRealVar("NLL", "-log(Likelihood)", 1.0)
ngen = RooRealVar("ngen", "number of generated events", self.options().nevents)
seed = RooRealVar("seed", "random seed", 0.0)
from ROOT import RooCategory
status = RooCategory("status", "fit status")
status.defineType("success", 0)
status.defineType("one", 1)
status.defineType("two", 2)
status.defineType("three", 3)
status.defineType("other", 4)
result_params.add(status)
result_params.add(NLL)
result_params.add(ngen)
result_params.add(seed)
# The dataset to store the results
from ROOT import RooDataSet
self._data = RooDataSet("result_data", "result_data", result_params)
data_params = self._data.get()
from ROOT import RooRandom
import struct, os
while self._data.numEntries() < self.options().ntoys:
# Get a good random seed, set it and store it
s = struct.unpack("I", os.urandom(4))[0]
RooRandom.randomGenerator().SetSeed(s)
seed.setVal(s)
# Reset pdf parameters to initial values. Note: this does not reset the estimated errors...
pdf_params.assignValueOnly(self.gen_params())
args = dict(NumEvents=self.options().nevents)
if "ProtoData" in kwargs:
args["ProtoData"] = kwargs.pop("ProtoData")
genPdf.getParameters(obs_set).assignValueOnly(gen_pdf_params)
data = genPdf.generate(obs_set, **args)
if transform:
data = transform(data)
if not data:
# Transform has failed
transform.set_params(data_params)
self._data.add(data_params)
continue
if data.isWeighted() and "SumW2Error" not in self.fit_opts():
self.fit_opts()["SumW2Error"] = False
j = 0
while j < 4:
fit_result = pdf.fitTo(data, NumCPU=self.options().ncpu, **(self.fit_opts()))
if fit_result.status() == 0:
fit_result.Print()
break
j += 1
if fit_result.status() != 0:
print "Fit result status = %s" % fit_result.status()
NLL.setVal(fit_result.minNll())
if fit_result.status() < 4:
status.setIndex(fit_result.status())
else:
status.setIndex(4)
for result_param in result_params:
data_param = data_params.find(result_param.GetName())
if isinstance(result_param, RooCategory):
data_param.setIndex(result_param.getIndex())
else:
data_param.setVal(result_param.getVal())
# This sets a symmetric error, but since we don't run Minos, that's ok
data_param.setError(result_param.getError())
if transform:
transform.set_params(data_params)
self._data.add(data_params)
return self.data()
