def __init__( self, **kwargs ) : Name = kwargs.pop('Name', 'timeResModelMG') namePF = self.getNamePrefix(kwargs) from P2VV.RooFitWrappers import ResolutionModel, AddModel from ROOT import RooNumber self._parseArg('time', kwargs, Title = 'Decay time', Unit = 'ps', Observable = True, Value = 0., MinMax = ( -0.5, 5. )) self._parseArg('sigmat', kwargs, Title = '#sigma(t)', Unit = 'ps', Observable = True, MinMax = (0.0001, 0.12) ) self._timeResMuSF = self._parseArg( 'timeResMuSF', kwargs, Value = 1.0, Constant = True ) self._timeResSigmaOffset = self._parseArg( 'timeResSigmaOffset', kwargs, Value = 0.01, Error = 0.001, MinMax = ( 0.00001, 1 ) ) sigmasSFs = kwargs.pop('ScaleFactors', [(2, 3), (1, 1)]) gexps = kwargs.pop('GExp', dict([(i[0], False) for i in sigmasSFs])) fracs = kwargs.pop('Fractions', [(2, 0.165)]) split_fracs = kwargs.pop('SplitFracs', True) self.__split_mean = kwargs.pop('SplitMean', False) self.__simultaneous = kwargs.pop('Simultaneous', False) self.__mu_param = kwargs.pop('MeanParameterisation', '') self._splitVars = [] sf_param = kwargs.pop('TimeResSFParam', False) from P2VV.RooFitWrappers import RealVar self.__sf_placeholder = self._parseArg('sf_placeholder', kwargs, Value = 0, MinMax = (-1e6, 1e6), Constant = True) self.__split_placeholders = kwargs.pop('SplitPlaceholders', False) if 'sigmat' in self.__mu_param or self.__split_placeholders: self.__mu_placeholder = self._parseArg('mu_placeholder', kwargs, Value = 0, MinMax = (-1e6, 1e6), Constant = True) else: self.__mu_placeholder = self.__sf_placeholder assert(len(sigmasSFs) - 1 == len(fracs)) assert(not (self.__simultaneous and sf_param)) ## If this is to be fitted in a non-simultaneous fit, the mean sigmat ## value should be passed in to get the parameterisation correct. assert(not (self.__split_mean and self.__mu_param)) assert(self.__mu_param in ['quadratic', 'linear', '', 'linear_sigmat', 'quadratic_sigmat']) self._timeResMu = self._parseArg( 'timeResMu', kwargs, Value = -0.0027, MinMax = ( -2, 2 ) ) if 'linear' in self.__mu_param: self._mu = self._timeResMu self._mu_offset = self._parseArg( 'timeResMu_offset', kwargs, Value = 0, MinMax = (-2, 2)) self._mu_slope = self._parseArg( 'timeResMu_slope', kwargs, Value = 0, MinMax = (-2, 2)) if not self.__simultaneous or 'sigmat' in self.__mu_param: formula = '@0 + @1 * (@2 - @3)' args = [self._mu_offset, self._mu_slope, self._sigmat, self.__mu_placeholder] else: formula = '@0 + @1 * @2' args = [self._mu_offset, self._mu_slope, self.__mu_placeholder] self._timeResMu = self._parseArg('timeResMu_linear', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args) elif 'quadratic' in self.__mu_param: self._mu = self._timeResMu self._mu_offset = self._parseArg( 'timeResMu_offset', kwargs, Value = -0.001723, MinMax = (-2, 2)) self._mu_slope = self._parseArg( 'timeResMu_slope', kwargs, Value = -0.2, MinMax = (-2, 2)) self._mu_quad = self._parseArg( 'timeResMu_quad', kwargs, Value = -1, MinMax = (-50, 50)) if not self.__simultaneous or 'sigmat' in self.__mu_param: formula = '@2 + @3 * (@0 - @1) + @4 * (@0 - @1) * (@0 - @1)' args = [self._sigmat, self.__mu_placeholder, self._mu_offset, self._mu_slope, self._mu_quad] else: formula = '@1 + @2 * @0 + @3 * @0 * @0' args = [self.__mu_placeholder, self._mu_offset, self._mu_slope, self._mu_quad] self._timeResMu = self._parseArg('timeResMu_quadratic', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args) self._cache = kwargs.pop('Cache', True) param = kwargs.pop('Parameterise', False) assert(param in [False, 'RMS', 'Comb']) self._timeResSigmasSFs = [ self._parseArg( 'timeResSigmaSF_%s' % num, kwargs, Value = val, MinMax = (0.001, 20) )\ for num, val in sigmasSFs ] self._timeResFracs = [ self._parseArg( 'timeResFrac%s' % num, kwargs, Value = val, MinMax = (0.01, 0.999))\ for num, val in fracs ] from ROOT import RooNumber RooInf = RooNumber.infinity() if param == 'RMS': if sf_param: self._parseArg('sf_mean_offset', kwargs, Value = 0.05, MinMax = (-0.1, 2.) ) self._parseArg('sf_mean_slope', kwargs, Value = 1.4, MinMax = (-10, 10) ) self._parseArg('sf_sigma_offset', kwargs, Value = 0.01, MinMax = (-0.1, 2.) ) self._parseArg('sf_sigma_slope', kwargs, Value = 0.4, MinMax = (-10, 10) ) if sf_param.startswith('quadratic'): self._parseArg('sf_mean_quad', kwargs, Value = 1, MinMax = (-20, 20)) self._parseArg('sf_sigma_quad', kwargs, Value = 1, MinMax = (-20, 20)) if sf_param == 'linear': formula = '@2 + @3 * (@0 - @1)' args = {'mean' : [self._sigmat, self.__sf_placeholder, self._sf_mean_offset, self._sf_mean_slope], 'sigma' : [self._sigmat, self.__sf_placeholder, self._sf_sigma_offset, self._sf_sigma_slope]} self._sf_mean = self._parseArg('timeResSFMean_linear', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args['mean']) self._sf_sigma = self._parseArg('timeResSFSigma_linear', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args['sigma']) elif sf_param == 'linear_no_offset': formula = '@1 * @0' args = {'mean' : [self._sigmat, self._sf_mean_slope], 'sigma' : [self._sigmat, self._sf_sigma_slope]} self._sf_mean = self._parseArg('timeResSFMean_linear', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args['mean']) self._sf_sigma = self._parseArg('timeResSFSigma_linear', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args['sigma']) elif sf_param == ('quadratic'): formula = '@2 + @3 * (@0 - @1) + @4 * (@0 - @1) * (@0 - @1)' args = {'mean' : [self._sigmat, self.__sf_placeholder, self._sf_mean_offset, self._sf_mean_slope, self._sf_mean_quad], 'sigma' : [self._sigmat, self.__sf_placeholder, self._sf_sigma_offset, self._sf_sigma_slope, self._sf_sigma_quad]} self._sf_mean = self._parseArg('timeResSFMean_quad', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args['mean']) self._sf_sigma = self._parseArg('timeResSFSigma_quad', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args['sigma']) elif sf_param == ('quadratic_no_offset'): formula = '(@2 + @3 * (@0 - @1)) * @0' args = {'mean' : [self._sigmat, self.__sf_placeholder, self._sf_mean_slope, self._sf_mean_quad], 'sigma' : [self._sigmat, self.__sf_placeholder, self._sf_sigma_slope, self._sf_sigma_quad]} self._sf_mean = self._parseArg('timeResSFMean_quad', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args['mean']) self._sf_sigma = self._parseArg('timeResSFSigma_quad', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args['sigma']) else: from math import sqrt self._sf_mean = self._parseArg('timeResSFMean', kwargs , Value = ((1. - fracs[-1][1]) * sigmasSFs[-1][1] + fracs[-1][1] * sigmasSFs[-2][1]) , MinMax = (0.001, 0.2 if self.__simultaneous else 5)) self._sf_sigma = self._parseArg('timeResSFSigma', kwargs, Value = sqrt((1 - fracs[-1][1]) * sigmasSFs[-1][1] * sigmasSFs[-1][1] + fracs[-1][1] * sigmasSFs[-2][1] * sigmasSFs[-2][1] - self._sf_mean.getVal() ** 2), MinMax = (0.001, 0.05 if self.__simultaneous else 5)) if self.__simultaneous: self._splitVars += [self._sf_mean, self._sf_sigma] self._timeResSigmasSFs[-1] = self._parseArg(Name + '_SF1', kwargs, Formula = '- sqrt(@0 / (1 - @0)) * @1 + @2', Arguments = (self._timeResFracs[-1], self._sf_sigma, self._sf_mean), ObjectType = 'FormulaVar') self._timeResSigmasSFs[-2] = self._parseArg(Name + '_SF2', kwargs, Formula = 'sqrt((1 - @0) / @0) * @1 + @2', Arguments = (self._timeResFracs[-1], self._sf_sigma, self._sf_mean), ObjectType = 'FormulaVar') else: if sf_param: self._parseArg('sf_one_offset', kwargs, Value = 0.4, MinMax = (-10, 10)) self._parseArg('sf_one_slope', kwargs, Value = 1., MinMax = (0.1, 20)) self._parseArg('sf_two_offset', kwargs, Value = 0.4, MinMax = (-10, 10)) self._parseArg('sf_two_slope', kwargs, Value = 1., MinMax = (0.1, 10)) if sf_param == 'linear': formula = '@2 + @3 * (@0 - @1)' args = {'one' : [self._sigmat, self.__sf_placeholder, self._sf_one_offset, self._sf_one_slope], 'two' : [self._sigmat, self.__sf_placeholder, self._sf_two_offset, self._sf_two_slope]} self._sf_one = self._parseArg('sf_one_linear', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args['one']) self._sf_two = self._parseArg('sf_two_linear', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args['two']) self._timeResSigmasSFs[-1] = self._sf_one self._timeResSigmasSFs[-2] = self._sf_two elif sf_param == 'linear_no_offset': formula = '@1 * @0' args = {'one' : [self._sigmat, self._sf_one_slope], 'two' : [self._sigmat, self._sf_two_slope]} self._sf_one = self._parseArg('sf_one_linear', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args['one']) self._sf_two = self._parseArg('sf_two_linear', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args['two']) self._timeResSigmasSFs[-1] = self._sf_one self._timeResSigmasSFs[-2] = self._sf_two elif sf_param == 'quadratic': self._parseArg('sf_one_quad', kwargs, Value = 0.1, MinMax = (-10, 10) ) self._parseArg('sf_two_quad', kwargs, Value = 0.1, MinMax = (-10, 10) ) formula = '@2 + @3 * (@0 - @1) + @4 * (@0 - @1) * (@0 - @1)' args = {'one' : [self._sigmat, self.__sf_placeholder, self._sf_one_offset, self._sf_one_slope, self._sf_one_quad], 'two' : [self._sigmat, self.__sf_placeholder, self._sf_two_offset, self._sf_two_slope, self._sf_two_quad]} self._sf_one = self._parseArg('timeResSFMean_quad', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args['one']) self._sf_two = self._parseArg('timeResSFSigma_quad', kwargs, Formula = formula, ObjectType = 'FormulaVar', Arguments = args['two']) self._timeResSigmasSFs[-1] = self._sf_one self._timeResSigmasSFs[-2] = self._sf_two elif self.__simultaneous: self._splitVars += self._timeResSigmasSFs if split_fracs: self._splitVars += self._timeResFracs self._check_extraneous_kw( kwargs ) from ROOT import RooGaussModel as GaussModel from ROOT import RooGExpModel as GExpModel models = [] for numVal, sigmaSF, in zip(sigmasSFs, self._timeResSigmasSFs): gexp = gexps[numVal[0]] if gexp: rlife = self._parseArg( 'rlife_%d' % numVal[0], kwargs, Value = 0.1, MinMax = ( 0.0001, 10 ) ) rlife_sf = self._parseArg( 'rlife_sf', kwargs, Value = 1, ObjectType = 'ConstVar' ) self._splitVars += [rlife] params = [self._time, self._timeResMu, sigmaSF, rlife, self._timeResMuSF, self._timeResMuSF, self._timeResMuSF, 'false', 'Normal'] else: params = [self._time, self._timeResMu, sigmaSF] model_type = GExpModel if gexp else GaussModel model = ResolutionModel( Name = '%stimeResGauss_%s' % (namePF, numVal[0]), Type = model_type, Parameters = params, ConditionalObservables = [self._sigmat]) models.append(model) TimeResolution.__init__(self, Name = Name , Model = AddModel(Name, Models = models, Fractions = self._timeResFracs) , Cache = self._cache)
obsKeys += [ 'wMC', 'sel', 'selClTail' , 'firstData', 'hlt2_prescale', 'polarity' , 'trigDecUnb', 'trigDecExclB' , 'tagDecOS', 'tagDecSS' ] if addMomenta : obsKeys += [ 'B_P', 'B_Pt', 'B_eta', 'B_phi' , 'Kplus_PX', 'Kplus_PY', 'Kplus_PZ', 'Kplus_LOKI_ETA' , 'Kminus_PX', 'Kminus_PY', 'Kminus_PZ', 'Kminus_LOKI_ETA' , 'muplus_PX', 'muplus_PY', 'muplus_PZ', 'muplus_LOKI_ETA' , 'muminus_PX', 'muminus_PY', 'muminus_PZ', 'muminus_LOKI_ETA' ] from math import pi from ROOT import RooNumber RooInf = RooNumber.infinity() obsDict = dict( runPeriod = ( 'runPeriod', 'run period', dict( p2011 = 2011, p2012 = 2012 ) ) , firstData = ( 'firstData', 'first data', dict( first = +1, later = 0 ) ) , hlt2_prescale = ( 'hlt2_prescale', 'HLT2 prescale', dict( presc = +1, noPresc = 0 ) ) , polarity = ( 'polarity', 'magnet polarity', dict( up = +1, down = -1, noPol = 0 ) ) , tagDecOS = ( 'tagdecision_os_cb', 'OS tag decision', dict( B = +1, Bbar = -1, Untagged = 0 ) ) , tagDecSS = ( 'tagdecision_ss_nn', 'SS tag decision', dict( B = +1, Bbar = -1, Untagged = 0 ) ) , sel = ( 'sel', 'selection', dict( sel = 1, notSel = 0 ) ) , selClTail = ( 'sel_cleantail', 'clean tail sel.', dict( sel = 1, notSel = 0 ) ) , hlt1ExclB = ( 'hlt1_excl_biased', 'HLT1 excl. B.', dict( exclB = 1, notExclB = 0 ) ) , hlt2B = ( 'hlt2_biased', 'HLT2 B.', dict( B = 1, notB = 0 ) ) , hlt2UB = ( 'hlt2_unbiased', 'HLT2 UB.', dict( UB = 1, notUB = 0 ) ) , trigDecUnb = ( 'triggerDecisionUnbiased', 'trigDecUnb', dict( UB = 1, notUB = 0 ) ) , trigDecExclB = ( 'triggerDecisionBiasedExcl', 'trigDecExclB', dict( exclB = 1, notExclB = 0 ) ) , wMC = ( 'wMC', 'pbkgWeight', '', 0., -RooInf, +RooInf ) , mass = ( 'mass', 'm(J/#psi K^{+}K^{-})', 'MeV/c^{2}', 5368., 5200., 5550. )
def reweigh(target, target_cat, source, source_cat, binning = None): cats = {} for cat, data, name in [(target_cat, target, 'target_cat'), (source_cat, source, 'source_cat')]: if hasattr(cat, '_target_'): cat = cat._target_() elif type(cat) == str: cat_name = cat cat = data.get().find(cat_name) if not cat: data.Print() raise RuntimeError('observable or category %s is not present in data' % cat_name) if not any([isinstance(cat, t) for t in [RooCategory, RooRealVar]]): raise RuntimeError('category must be either a RooRealVar or a RooCategory') if isinstance(cat, RooRealVar): assert(binning) if type(binning) == array: binning = RooBinning(len(binning) - 1, binning) binning.SetName('reweigh') cat.setBinning(binning, 'reweigh') cat_name = cat.GetName() + '_cat_' + data.GetName() test_cat = data.get().find(cat_name) if not test_cat: cat = BinningCategory(Name = cat_name, Observable = cat, Binning = binning, Data = data, Fundamental = True) else: cat = test_cat cats[name] = cat target_cat = cats['target_cat'] source_cat = cats['source_cat'] print target_cat print source_cat target_bins = dict([(ct.getVal(), ct.GetName()) for ct in target_cat]) source_bins = dict([(ct.getVal(), ct.GetName()) for ct in source_cat]) print target_bins print source_bins target_table = target.table(target_cat) source_table = source.table(source_cat) target_table.Print('v') source_table.Print('v') from collections import defaultdict reweigh_weights = {} for i, l in sorted(target_bins.iteritems()): sf = source_table.getFrac(source_bins[i]) if sf == 0: w = 0 else: try: w = sf / target_table.getFrac(l) except ZeroDivisionError: print 'Warning bin %s in wpv_data is 0, setting weight to 0' % l w = 0. reweigh_weights[i] = w # RooFit infinity from ROOT import RooNumber RooInf = RooNumber.infinity() from RooFitWrappers import RealVar weight_var = RealVar('reweigh_var', MinMax = (RooInf, RooInf)) from ROOT import RooDataSet data_name = target.GetName() + 'weight_data' from ROOT import RooArgSet weight_data = RooDataSet(data_name, data_name, RooArgSet(weight_var)) weight_var = weight_data.get().find(weight_var.GetName()) for i in range(target.numEntries()): r = target.get(i) n = target_cat.getIndex() w = reweigh_weights[n] weight_var.setVal(target.weight() * w) weight_data.fill() target.merge(weight_data) target = RooDataSet(target.GetName(), target.GetTitle(), target, target.get(), '', weight_var.GetName()) return target, reweigh_weights