total_entries = 450 basevars = BaseVariables(0, 0.1444, 0.5, 3.5) basevars.get_time().setConstant(True) # Set up the WIMP class wimp_class = WIMPModel(basevars, mass_of_wimp=wimp_mass, kilograms=0.4, constant_quenching=False) model = wimp_class.get_model() # Set up the background class low_energy = LowEnergyBackgroundModel(basevars) low_energy_model = low_energy.get_model() list_of_models, list_of_coefficients = low_energy.get_list_components() background_normal = ROOT.RooRealVar("flat_normal", "Background event number", 0, total_entries) background_extend = ROOT.RooExtendPdf("background_extend", "background_extend", low_energy_model, background_normal) # Now set up the extended model model_normal = ROOT.RooRealVar("model_normal", "WIMP-nucleus xs", 1, -10, total_entries, "pb") model_extend = ROOT.RooExtendPdf("model_extend", "model_extend", model, model_normal) variables = ROOT.RooArgSet(basevars.get_energy()) # Following is the number of events # Remember model_normal is in pb, so this is
class DataExclusion(WIMPModel): def get_requested_values(cls): adict = WIMPModel.get_requested_values() del adict['constant_energy'] del adict['constant_time'] del adict['background_rate'] adict['fix_l_line_ratio'] = ('Fix ratio of the Ge and Zn L-lines', False) adict['data_file'] = ('Name of data root file', 'temp.root') adict['object_name'] = ("""Name of object inside data file. This can be a: TH1: the bins will be interprected as energy. A RooDataHist will be used to generate a binned fit. TTree: with branches, 'ee_energy', 'time', 'weight', plus others. 'time' and ee_energy are both optional, it is assumed that if they don't exist that time and energy are constant. 'weight' is optional, but when present will adjust the weight of each entry to 'weight' in the RooDataSet. A RooDataSet will be used to generate an unbinned fit. """, 'output_data') adict['data_set_cuts'] = ("""String of cuts to apply. This only applies if the object passed in is a TTree (see object_name). These cuts will be used to generate a subset of the TTree and pass into RooDataSet. """, '') return adict get_requested_values = classmethod(get_requested_values) def initialize(self): # The background model is the same, but now we switch it to the data model from pyWIMP.DMModels.low_energy_background import LowEnergyBackgroundModel from pyWIMP.DMModels.base_model import BaseVariables from pyWIMP.DMModels.wimp_model import WIMPModel self.total_counts = None open_file = ROOT.TFile(self.data_file) self.workspace = open_file.Get(self.object_name) # Do some introspection, we can handle TH1s, and RooAbsData self.basevars = BaseVariables(time_beginning=0, time_in_years=self.total_time, energy_threshold=self.threshold, energy_max=self.energy_max, use_tag=False) self.variables = ROOT.RooArgSet() if self.workspace.InheritsFrom(ROOT.TH1.Class()): self.variables.add(self.basevars.get_energy()) self.basevars.get_time().setVal(0) self.basevars.get_time().setConstant(True) self.data_set_model = ROOT.RooDataHist("data", "data", ROOT.RooArgList(self.basevars.get_energy()), self.workspace) elif self.workspace.InheritsFrom(ROOT.TTree.Class()): # Default to setting them to constant. self.basevars.get_time().setVal(0) self.basevars.get_time().setConstant(True) self.basevars.get_energy().setVal(0) self.basevars.get_energy().setConstant(True) if self.num_energy_bins != 0: self.basevars.get_energy().setBins(int(self.num_energy_bins)) if self.num_time_bins != 0: self.basevars.get_time().setBins(int(self.num_time_bins)) branches = self.workspace.GetListOfBranches() efficiency = "" branch_arg_list = [] for i in range(branches.GetEntries()): branch_name = branches[i].GetName() if branch_name == "ee_energy": self.basevars.get_energy().setConstant(False) self.variables.add(self.basevars.get_energy()) branch_arg_list.append((self.basevars.get_energy(), branch_name)) if branch_name == "time": self.basevars.get_time().setConstant(False) self.variables.add(self.basevars.get_time()) branch_arg_list.append((self.basevars.get_time(), branch_name)) elif branch_name == "weight": efficiency = branch_name self.variables.add(self.basevars.get_weighting()) if not self.data_set_cuts: # Load the DataSet the easy way self.data_set_model = ROOT.RooDataSet("data", "data", self.workspace, self.variables, efficiency) else: # Otherwise, we have to get the correct events, # which requires stepping through all events self.data_set_model = ROOT.RooDataSet("data", "data", self.variables, efficiency) # Get an event list with the correct cut events ROOT.gROOT.cd() el = ROOT.TEventList("el", "el") cuts = self.data_set_cuts iter = self.variables.createIterator() while 1: obj = iter.Next() if not obj: break if obj.GetTitle() == 'weight': continue cuts += " && ((%s <= %f) && (%s >= %f))" % (obj.GetName(), obj.getMax(), obj.GetName(), obj.getMin()) self.workspace.Draw(">>%s" % el.GetName(), cuts, "goff") event_list = [el.GetEntry(i) for i in range(el.GetN())] #event_list = [i for i in range(self.workspace.GetEntries())] for j in event_list: self.workspace.GetEntry(j) eff_val = 1. if efficiency: eff_val = getattr(self.workspace, efficiency) for val in branch_arg_list: val[0].setVal(getattr(self.workspace, val[1])) if eff_val != 0: self.data_set_model.add(self.variables, eff_val) else: print "Requested: %s, isn't a TTree or TH1!" % self.data_set_name raise TypeError if self.num_energy_bins != 0 or self.num_time_bins != 0: self.original_data_set = self.data_set_model self.data_set_model = self.original_data_set.binnedClone() if self.data_set_model.isWeighted(): print "Data set is weighted" print "Data set has %i entries." % self.data_set_model.sumEntries() if not self.do_axioelectric: self.wimpClass = WIMPModel(self.basevars, mass_of_wimp=self.wimp_mass, kilograms = self.mass_of_detector, constant_quenching=(not self.variable_quenching)) # This is where we define our model self.model = self.wimpClass.get_model() #self.model = self.wimpClass.get_simple_model() self.norm = self.wimpClass.get_normalization().getVal() # The following has not been normalized to per-nucleon yet. self.model_normal = ROOT.RooRealVar("model_normal", "WIMP-nucleus #sigma", 1, -10, 100, "pb") self.model_extend = ROOT.RooExtendPdf("model_extend", "model_extend", self.model, self.model_normal) # Getting the number of events for a model_normal of 1 # This gives us number of events per model_normal value scaler = self.model_extend.expectedEvents(self.variables) self.model_normal.setMax(2*self.data_set_model.sumEntries()/scaler) else: # wimpClass is of course a misnomer here, but we use it for now self.wimpClass = GaussianSignalModel(self.basevars, mean_of_signal=self.axion_mass) self.model_normal = ROOT.RooRealVar("model_normal", "Counts", 0, -10, 1000) # This is where we define our model self.model = self.wimpClass.get_model() self.norm = self.wimpClass.get_normalization()*self.model.getNorm( ROOT.RooArgSet(self.basevars.get_energy())) # We actually want the inverse of the normaliation, since this is later multiplied # and we want the total counts in a particular gaussian. self.norm = 1./self.norm self.model_extend = ROOT.RooExtendPdf("model_extend", "model_extend", self.model, self.model_normal) self.is_initialized = True self.calculation_class = \ dat.DataCalculation(self.exit_manager) self.low_energy = LowEnergyBackgroundModel(self.basevars, use_ratio=self.fix_l_line_ratio) list_of_models, list_of_coefficients = self.low_energy.get_list_components() self.extended_models = [] i = 0 while 1: amod = list_of_models.at(i) avar = list_of_coefficients.at(i) if not amod: break i += 1 extend = ROOT.RooExtendPdf("extend%s" % amod.GetName(), "extend%s" % amod.GetName(), amod, avar) self.extended_models.append(extend) temp_list = ROOT.RooArgList() temp_list.add(self.model_extend) for amod in self.extended_models: temp_list.add(amod) self.added_pdf = ROOT.RooAddPdf("b+s", "Background + Signal", temp_list) self.test_variable = self.model_normal self.fitting_model = self.added_pdf
basevars = BaseVariables(0, 0.14887063655,0.5, 3.5) basevars.get_time().setConstant(True) # Set up the WIMP class wimp_class = WIMPModel(basevars, mass_of_wimp = wimp_mass, kilograms=0.4, constant_quenching=False) model = wimp_class.get_model() normalization = wimp_class.get_normalization().getVal() # Set up the background class low_energy = LowEnergyBackgroundModel(basevars) low_energy_model = low_energy.get_model() list_of_models, list_of_coefficients = low_energy.get_list_components() background_normal = ROOT.RooRealVar("flat_normal", "Background event number", 0, 1000) background_extend = ROOT.RooExtendPdf("background_extend", "background_extend", low_energy_model, background_normal) #flat_coef.setMin(-5) #exp_coef.setMin(-5) # Now set up the extended model model_normal = ROOT.RooRealVar("model_normal", "WIMP-nucleus xs",