def fillIntoTree(out_tree, branches, cfg, hist_cfg, vcfgs, total_scale, plot, verbose, friend_func): if isinstance(cfg, HistogramCfg): # Loop over sub-cfgs and fill them total_scale *= cfg.total_scale if cfg.total_scale else 1. for sub_cfg in cfg.cfgs: fillIntoTree(out_tree, branches, sub_cfg, cfg, vcfgs, total_scale, plot, verbose, friend_func) return file_name = '/'.join([cfg.ana_dir, cfg.dir_name, cfg.tree_prod_name, 'tree.root']) # Attaches tree to plot ttree = plot.readTree(file_name, cfg.tree_name, verbose=verbose, friend_func=friend_func) norm_cut = hist_cfg.cut shape_cut = hist_cfg.cut if cfg.norm_cut: norm_cut = cfg.norm_cut if cfg.shape_cut: shape_cut = cfg.shape_cut full_weight = branches[-1] weight = hist_cfg.weight if cfg.weight_expr: weight = '*'.join([weight, cfg.weight_expr]) if hist_cfg.weight: norm_cut = '({c}) * {we}'.format(c=norm_cut, we=weight) shape_cut = '({c}) * {we}'.format(c=shape_cut, we=weight) # and this one too sample_weight = cfg.scale * total_scale if not cfg.is_data: sample_weight *= hist_cfg.lumi*cfg.xsec/cfg.sumweights formula = TTreeFormula('weight_formula', norm_cut, ttree) formula.GetNdata() # Add weight as tree variable # Then loop over ttree # And save this to the other tree # # Create TTreeFormulas for all vars for var in vcfgs: if var.drawname != var.name: var.formula = TTreeFormula('formula'+var.name, var.drawname, ttree) var.formula.GetNdata() for i in xrange(ttree.GetEntries()): ttree.GetEntry(i) w = formula.EvalInstance() if w == 0.: continue full_weight[0] = w * sample_weight if abs(full_weight[0]) > 1000.: print "WARNING, unusually large weight", w, sample_weight import pdb; pdb.set_trace() print '\nWeight:', full_weight[0] print cfg.name print norm_cut for branch, var in zip(branches, vcfgs): branch[0] = var.formula.EvalInstance() if hasattr(var, 'formula') else getattr(ttree, var.name) out_tree.Fill() if shape_cut != norm_cut: print 'WARNING: different norm and shape cuts currently not supported in HistCreator.createTrees'
tree_out = tree_in.CopyTree(cut_str) tree_out.Write() # file_out.Close() new_file_out = TFile(file_out_name.replace('.root', '_weight.root'), 'RECREATE') weight_tree = tree_out.CloneTree(0) scale = int_lumi * sample.xsec * sample.scale / sample.sumweights full_weight = array('f', [0.]) new_b = weight_tree.Branch('full_weight', full_weight, 'full_weight/F') formula = TTreeFormula('weight_formula', weight, tree_out) formula.GetNdata() # ATTENTION THIS MAY NOT WORK! for i in xrange(tree_out.GetEntries()): tree_out.GetEntry(i) full_weight[0] = formula.EvalInstance() * scale # print full_weight[0] # new_b.Fill() weight_tree.Fill() # tree_out.Fill() new_file_out.Write() new_file_out.Close() file_out.Close() print 'Writing file', file_out_name