def make_staged_target(stage_dir, target): if not isinstance(stage_dir, LocalDirectoryTarget): stage_dir = LocalDirectoryTarget(stage_dir) return stage_dir.child(target.unique_basename, type=target.type, **target._copy_kwargs())
def prepare_stageout(self): outputs = mask_struct(self.task.sandbox_stageout_mask(), self.task.output()) if not outputs: return None # create a tmp dir tmp_dir = LocalDirectoryTarget(is_tmp=True) tmp_dir.touch() # map output files to local local targets in tmp_dir def map_target(target): return make_staged_target(tmp_dir, target) stage_outputs = map_struct(map_target, outputs) return StageInfo(outputs, tmp_dir, stage_outputs)
def stagein(self): inputs = mask_struct(self.task.sandbox_stagein_mask(), self.task.input()) if not inputs: return None # create a tmp dir tmp_dir = LocalDirectoryTarget(is_tmp=True) tmp_dir.touch() # copy input files and map to local targets in tmp_dir def map_target(target): tmp_target = make_staged_target(tmp_dir, target) target.copy(tmp_target) return tmp_target stage_inputs = map_struct(map_target, inputs) return StageInfo(inputs, tmp_dir, stage_inputs)
def from_directory(cls, directory, **kwargs): # dir should be a FileSystemDirectoryTarget or a string, in which case it is interpreted as # a local path if isinstance(directory, six.string_types): d = LocalDirectoryTarget(directory) elif isinstance(d, FileSystemDirectoryTarget): d = directory else: raise TypeError("directory must either be a string or a FileSystemDirectoryTarget " "object, got '{}'".format(directory)) # find all files, pass kwargs which may filter the result further kwargs["type"] = "f" basenames = d.listdir(**kwargs) # convert to file targets targets = [d.child(basename, type="f") for basename in basenames] return cls(targets)
def localize(self, *args, **kwargs): # when localizing collections using temporary files, it makes sense to put # them all in the same temporary directory tmp_dir = kwargs.get("tmp_dir") if not tmp_dir: tmp_dir = LocalDirectoryTarget(is_tmp=True) kwargs["tmp_dir"] = tmp_dir # enter localize contexts of all targets with localize_file_targets(self.targets, *args, **kwargs) as localized_targets: # create a copy of this collection that wraps the localized targets yield self.__class__(localized_targets, **self._copy_kwargs())
def run(self): # before_run hook if callable(self.task.sandbox_before_run): self.task.sandbox_before_run() # create a temporary direction for file staging tmp_dir = LocalDirectoryTarget(is_tmp=True) tmp_dir.touch() # stage-in input files stagein_info = self.stagein(tmp_dir) if stagein_info: # tell the sandbox self.sandbox_inst.stagein_info = stagein_info logger.debug("configured sandbox stage-in data") # prepare stage-out stageout_info = self.prepare_stageout(tmp_dir) if stageout_info: # tell the sandbox self.sandbox_inst.stageout_info = stageout_info logger.debug("configured sandbox stage-out data") # create the actual command to run cmd = self.sandbox_inst.cmd(self.proxy_cmd()) # run with log section before and after actual run call with self._run_log(cmd): code, out, err = self.sandbox_inst.run(cmd) if code != 0: raise Exception("Sandbox '{}' failed with exit code {}".format( self.sandbox_inst.key, code)) # actual stage_out if stageout_info: self.stageout(stageout_info) # after_run hook if callable(self.task.sandbox_after_run): self.task.sandbox_after_run()
def localize(self, path=None, is_tmp=True, skip_parent=False, mode=0o0770, cache=True, retry=None): tmp = LocalDirectoryTarget(path, is_tmp=is_tmp) try: yield tmp if not skip_parent: self.parent.touch() # determine files to put lpaths, rpaths = [], [] for base, dirs, files in tmp.walk(): for f in files: full_f = os.path.join(base, f) lpaths.append(full_f) rpaths.append(os.path.relpath(full_f, tmp.path)) self.put(lpaths, rpaths, cache=cache, retry=retry) self.chmod(mode, retry=retry) finally: del tmp
def localize(self, *args, **kwargs): # when localizing collections using temporary files, it makes sense to put # them all in the same temporary directory tmp_dir = kwargs.get("tmp_dir") tmp_dir_created = False if not tmp_dir: tmp_dir = LocalDirectoryTarget(is_tmp=True) tmp_dir_created = True kwargs["tmp_dir"] = tmp_dir.path # enter localize contexts of all targets with localize_file_targets(self.targets, *args, **kwargs) as localized_targets: # create a copy of this collection that wraps the localized targets copy = self.__class__(localized_targets, **self._copy_kwargs()) try: yield copy finally: # although tmp_dir would clean itself during garbage collection, an error might have # occurred, so for larger collections it is safer to delete the tmp_dir manually if tmp_dir_created: tmp_dir.remove()
def run(self): import ROOT ROOT.PyConfig.IgnoreCommandLineOptions = True ROOT.gROOT.SetBatch() inp = self.input() outp = self.output() local_tmp = LocalDirectoryTarget(is_tmp=True) local_tmp.touch() plots = {} if self.norm_to_nominal and self.shifts[0] != "nominal": raise KeyError("'norm_to_nominal' is set to true, but no nominal values found.") for color_idx, (config, config_input) in enumerate(inp.items()): b_tagger, iteration, version = config config_ids = [b_tagger] if len(self.iterations) > 1: config_ids.append("iteration {}".format(iteration)) if len(self.versions) > 1: config_ids.append("version {}".format(version)) config_id = ", ".join(config_ids) nominal_hists = {} nominal_fit_hists = {} # combined errors for multiple shifts up_shifted_hists = defaultdict(dict) up_shifted_fit_hists = defaultdict(dict) down_shifted_hists = defaultdict(dict) down_shifted_fit_hists = defaultdict(dict) if self.fix_normalization: normalization_input = config_input.pop("norm") for shift_idx, (shift, inp_target) in enumerate(config_input.items()): # get scaling factors for normalization if self.fix_normalization: norm_factors = normalization_input.load()[shift] with inp_target["fit"]["sf"].load("r") as fit_file, \ inp_target["hist"]["scale_factors"].load("r") as hist_file: for category_key in fit_file.GetListOfKeys(): category_name = category_key.GetName() if not self.config_inst.has_category(category_name): raise KeyError("Unknown category {}".format(category_name)) category = self.config_inst.get_category(category_name) pt_range = category.get_aux("pt") eta_range = category.get_aux("eta") region = category.get_aux("region") # same category name for different b-taggers if len(self.b_taggers) > 1: plot_category = category.name.replace("__" + b_tagger, "") else: plot_category = category.name fit_category_dir = fit_file.Get(category_name) fit_hist = fit_category_dir.Get(self.hist_name) hist_category_dir = hist_file.Get(category_name) hist = hist_category_dir.Get(self.hist_name) # truncate first bin hist = self.rebin_hist(hist, region, b_tagger=b_tagger, truncate=True) # normalize histogram if required # fit histograms are already normalized in FitScaleFactors if self.fix_normalization and not self.is_c_flavour: hist.Scale(norm_factors[category_name]) if shift == "nominal": # make sure histograms are not cleaned up when the file is closed nominal_fit_hists[plot_category] = fit_hist.Clone() nominal_fit_hists[plot_category].SetDirectory(0) nominal_hists[plot_category] = hist.Clone() nominal_hists[plot_category].SetDirectory(0) # for c-jets, there is no nominal histogram # Instead, all nominal values are set to 1 if shift_idx == 0 and self.is_c_flavour: nominal_fit_hist = fit_hist.Clone() for bin_idx in range(1, nominal_fit_hist.GetNbinsX() + 1): nominal_fit_hist.SetBinContent(bin_idx, 1.0) nominal_fit_hist.SetDirectory(0) nominal_fit_hists[plot_category] = nominal_fit_hist if shift != "nominal" and self.multiple_shifts: # collect all shifted fit histograms to build envelope later sys, direction = shift.rsplit("_", 1) if direction == "up": up_shifted_fit_hists[plot_category][sys] = fit_hist.Clone() up_shifted_fit_hists[plot_category][sys].SetDirectory(0) up_shifted_hists[plot_category][sys] = hist.Clone() up_shifted_hists[plot_category][sys].SetDirectory(0) elif direction == "down": down_shifted_fit_hists[plot_category][sys] = fit_hist.Clone() down_shifted_fit_hists[plot_category][sys].SetDirectory(0) down_shifted_hists[plot_category][sys] = hist.Clone() down_shifted_hists[plot_category][sys].SetDirectory(0) else: raise ValueError("Unknown direction {}".format(direction)) if self.norm_to_nominal: fit_hist.Divide(nominal_fit_hists[plot_category]) # get same category key for all b-taggers if plot_category in plots: plot = plots[plot_category] else: plot = ROOTPlot(category.name, category.name) plot.create_pads() plots[plot_category] = plot plot.cd(0, 0) fit_hist.GetXaxis().SetRangeUser(-.1, 1.0) y_min = 0.6 if self.norm_to_nominal else 0. y_max = 1.4 if self.norm_to_nominal else 2. fit_hist.GetYaxis().SetRangeUser(y_min, y_max) if len(self.b_taggers) == 1: title = self.config_inst.get_aux("btaggers")[b_tagger]["label"] + " discriminator" else: title = "B-Tag Discriminant" fit_hist.GetXaxis().SetTitle(title) fit_hist.GetYaxis().SetTitle("SF") if shift_idx == 0: if not self.multiple_shifts or shift == "nominal": # only draw this fit histogram if it is not part of a shifted envelope plot.draw({"sf": fit_hist}, line_color=1, add_to_legend=False) line = ROOT.TLine(0., 0., 0., 2.) line.SetLineStyle(9) plot.draw({"line": line}, add_same_option=False, line_color=1, add_to_legend=False) # add category information to plot if not np.isinf(pt_range[1]): text = r"#splitline{%d < p_{T} < %d}{%.1f < |#eta| < %.1f}" % \ (pt_range[0], pt_range[1], eta_range[0], eta_range[1]) else: text = r"#splitline{p_{T} > %d}{%.1f < |#eta| < %.1f}" % \ (pt_range[0], eta_range[0], eta_range[1]) plot.draw_text(text) elif not self.multiple_shifts: plot.draw({shift: fit_hist}, line_color=None) if shift == "nominal" and not self.norm_to_nominal: plot.draw({config_id + ", nominal": hist}, line_color=1, add_to_legend=(len(self.shifts) != 1)) if self.multiple_shifts: for plot_category in plots: plot = plots[plot_category] plot.cd(0, 0) # build shifted histograms fit_hist_down, fit_hist_up = build_hist_envelope(nominal_fit_hists[plot_category], up_shifted_fit_hists[plot_category], down_shifted_fit_hists[plot_category], envelope_as_errors=False) hist_down, hist_up = build_hist_envelope(nominal_hists[plot_category], up_shifted_hists[plot_category], down_shifted_hists[plot_category], envelope_as_errors=False) if self.norm_to_nominal: fit_hist_up.Divide(nominal_fit_hists[plot_category]) fit_hist_down.Divide(nominal_fit_hists[plot_category]) hist_up.Divide(nominal_hists[plot_category]) hist_down.Divide(nominal_hists[plot_category]) plot.draw({config_id + ", up": fit_hist_up}, line_color=None) plot.draw({config_id + ", down": fit_hist_down}, line_color=None) plot.draw({config_id + ", up": hist_up}, line_color=2, options=["hist"]) plot.draw({config_id + ", down": hist_down}, line_color=4, options=["hist"]) # save plots for plot_category in plots: plot = plots[plot_category] plot_name = self.get_plot_name(plot_category, self.shifts_identifier, self.b_taggers[0], self.iterations[0]) plot.save(os.path.join(local_tmp.path, plot_name), draw_legend=True, lumi=self.config_inst.get_aux("lumi").values()[0] / 1000.) del plot with outp.localize("w") as tmp: with tarfile.open(tmp.path, "w:gz") as tar: for plot_file in os.listdir(local_tmp.path): tar.add(os.path.join(local_tmp.path, plot_file), arcname=plot_file)
def env(self): # strategy: unlike docker, singularity might not allow binding of paths that do not exist # in the container, so create a tmp directory on the host system and bind it as /tmp, let # python dump its full env into a file, and read the file again on the host system if self.image not in self._envs: tmp_dir = LocalDirectoryTarget(is_tmp=True) tmp_dir.touch() tmp = tmp_dir.child("env", type="f") tmp.touch() # determine whether volume binding is allowed allow_binds_cb = getattr(self.task, "singularity_allow_binds", None) if callable(allow_binds_cb): allow_binds = allow_binds_cb() else: cfg = Config.instance() allow_binds = cfg.get_expanded(self.get_config_section(), "allow_binds") # arguments to configure the environment args = ["-e"] if allow_binds: args.extend(["-B", "{}:/tmp".format(tmp_dir.path)]) env_file = "/tmp/{}".format(tmp.basename) else: env_file = tmp.path # get the singularity exec command singularity_exec_cmd = self._singularity_exec_cmd() + args # build commands to setup the environment setup_cmds = self._build_setup_cmds(self._get_env()) # build the python command that dumps the environment py_cmd = "import os,pickle;" \ + "pickle.dump(dict(os.environ),open('{}','wb'),protocol=2)".format(env_file) # build the full command cmd = quote_cmd(singularity_exec_cmd + [ self.image, "bash", "-l", "-c", "; ".join( flatten(setup_cmds, quote_cmd(["python", "-c", py_cmd]))), ]) # run it code, out, _ = interruptable_popen(cmd, shell=True, executable="/bin/bash", stdout=subprocess.PIPE, stderr=subprocess.STDOUT) if code != 0: raise Exception( "singularity sandbox env loading failed:\n{}".format(out)) # load the environment from the tmp file env = tmp.load(formatter="pickle") # cache self._envs[self.image] = env return self._envs[self.image]
def run(self): ROOT.PyConfig.IgnoreCommandLineOptions = True ROOT.gROOT.SetBatch() ROOT.gSystem.Load('libCondFormatsBTauObjects') ROOT.gSystem.Load('libCondToolsBTau') local_tmp = LocalDirectoryTarget(is_tmp=True) local_tmp.touch() jes_sources = self.config_inst.get_aux("jes_sources") shifts = [] if self.shift == "ALL": if self.flavor == "c": shifts.extend(["cferr1", "cferr2"]) else: shifts.extend([ "jes{}".format(jes_source) for jes_source in jes_sources if jes_source != "Total" ]) shifts.extend([ "{}{}".format(region, type) for region, type in itertools.product(["lf", "hf"], ["", "stats1", "stats2"]) ]) elif self.shift == "NONE": shifts = [] else: shifts = [self.shift] v_sys = getattr(ROOT, 'vector<string>')() for shift in shifts: v_sys.push_back("up_" + shift) v_sys.push_back("down_" + shift) flavor_ids = self.config_inst.get_aux("flavor_ids") binning = self.config_inst.get_aux("binning")[self.flavor] pt_binning = [ (start, end) for start, end in zip(binning["pt"][:-1], binning["pt"][1:]) ] eta_binning = [(start, end) for start, end in zip( binning["abs(eta)"][:-1], binning["abs(eta)"][1:])] figures = {} if self.compare_file is None: csv_files, descriptions = [self.csv_file], ["csv"] else: csv_files, descriptions = [self.csv_file, self.compare_file], ["new", "old"] for input_file, id in zip(csv_files, descriptions): # create calibration reader calib = ROOT.BTagCalibration("csv_{}".format(id), input_file) reader = ROOT.BTagCalibrationReader( 3, # 0 is for loose op, 1: medium, 2: tight, 3: discr. reshaping "central", # central systematic type v_sys, # vector of other sys. types ) for jetFlavor in [0, 1, 2]: reader.load( calib, jetFlavor, # 0 is for b flavour, 1: FLAV_C, 2: FLAV_UDSG "iterativefit" # measurement type ) for pt_idx, pt_range in enumerate(pt_binning): for eta_idx, eta_range in enumerate(eta_binning): key = pt_range + eta_range if key not in figures: fig = plt.figure() ax = fig.add_subplot(111) ax.set_title("pt: %s to %s, eta: %.1f to %.1f" % (pt_range + eta_range)) else: fig, ax = figures[key] if pt_range[1] == np.inf: pt_val = pt_range[0] + 1 else: pt_val = np.mean(pt_range) eta_val = np.mean(eta_range) def get_values(csv_reader, sys_type): x_values = np.linspace(-0.1, 1., 10000) y_values = [] for csv_value in x_values: sf = csv_reader.eval_auto_bounds( sys_type, # systematic (here also 'up'/'down' possible) flavor_ids[self.flavor], # jet flavor eta_val, # absolute value of eta pt_val, # pt csv_value) y_values.append(sf) return np.array(x_values), np.array(y_values) x_values, nominal_values = get_values(reader, "central") if not self.norm_to_nominal: ax.plot(x_values, nominal_values, label="{}, {}".format(id, "nominal")) if self.shift != "NONE": total_errors_up = np.zeros(nominal_values.shape) total_errors_down = np.zeros(nominal_values.shape) for shift in shifts: _, up_values = get_values(reader, "up_" + shift) _, down_values = get_values( reader, "down_" + shift) if len(shifts) > 1: # build envelope diff_up = up_values - nominal_values diff_down = down_values - nominal_values # shift with effect in up/down direction errors_up = np.max([ diff_up, diff_down, np.zeros(nominal_values.shape) ], axis=0) errors_down = np.min([ diff_up, diff_down, np.zeros(nominal_values.shape) ], axis=0) # add in quadrature total_errors_up += errors_up**2 total_errors_down += errors_down**2 total_errors_up = total_errors_up**0.5 total_errors_down = total_errors_down**0.5 if len(shifts) > 1: up_values = nominal_values + total_errors_up down_values = nominal_values - total_errors_down if self.norm_to_nominal: up_values /= nominal_values down_values /= nominal_values ax.plot(x_values, up_values, label="{}, {}".format(id, "up_" + self.shift)) ax.plot(x_values, down_values, label="{}, {}".format(id, "down" + self.shift)) if self.compare_file is None: if self.flavor in ["c", "hf"]: input_file = self.root_hf_file elif self.flavor == "lf": input_file = self.root_lf_file else: raise Exception("No .root file for c flavor SFs.") root_file = ROOT.TFile.Open(input_file) func_name = "csv_ratio_Pt{}_Eta{}_final".format( pt_idx, eta_idx) if self.flavor == "c": func_name = "c_" + func_name func = root_file.Get(func_name) y_values = [] x_values = np.linspace(-0.1, 1., 10000) for csv_value in x_values: y_val = func.Eval(csv_value) y_values.append(y_val) ax.plot(x_values, y_values, label="{}, {}".format(".root", "nominal")) figures[key] = (fig, ax) del reader del calib for key, (fig, ax) in figures.items(): ax.legend(loc="lower right") ax.set_ylim(0., 2.) fig.savefig( os.path.join( local_tmp.path, "SF_%s_%s_Pt%sTo%s_eta%.1fTo%.1f.pdf" % ((self.flavor, self.shift) + key))) with self.output().localize("w") as tmp: with tarfile.open(tmp.path, "w:gz") as tar: for plot_file in os.listdir(local_tmp.path): tar.add(os.path.join(local_tmp.path, plot_file), arcname=plot_file)
def run(self): ROOT.PyConfig.IgnoreCommandLineOptions = True ROOT.gROOT.SetBatch() local_tmp = LocalDirectoryTarget(is_tmp=True) local_tmp.touch() if self.flavor == "hf": input_file = self.hf_file else: input_file = self.lf_file root_file = ROOT.TFile.Open(input_file) hist_tpl = "h_csv_ratio_Pt{}_Eta{}_final" binning = self.config_inst.get_aux("binning")[self.flavor] n_pt_categories = len(binning["pt"]) - 1 n_eta_categories = len(binning["abs(eta)"]) - 1 hist_names = [ hist_tpl.format(pt_idx, eta_idx) for pt_idx in range(n_pt_categories) for eta_idx in range(n_eta_categories) ] ## ## Plot histograms with total systematic envelope ## for hist_name in hist_names: nominal_hist = root_file.Get(hist_name) errors_up = [] errors_down = [] # collect all shifts from file if self.shift == "ALL": shifts = [ key.GetName().split("_")[-1] for key in root_file.GetListOfKeys() if key.GetName().startswith(hist_name) ] shifts = set( [shift[:-2] for shift in shifts if shift[-2:] == "Up"]) else: shifts = [self.shift] for shift_idx, shift in enumerate(shifts): hist_name_up = hist_name + "_" + shift + "Up" hist_name_down = hist_name + "_" + shift + "Down" hist_up = root_file.Get(hist_name_up) hist_down = root_file.Get(hist_name_down) for bin_idx in range(1, nominal_hist.GetNbinsX() + 1): nominal_value = nominal_hist.GetBinContent(bin_idx) # combine all shifts that have an effect in the same direction # effect from <shift>_up/done systematics diff_up = hist_up.GetBinContent(bin_idx) - nominal_value diff_down = hist_down.GetBinContent( bin_idx) - nominal_value # detect systematics where up/down shift direction is the same #if diff_up * diff_down > 0: # print "One sided shift: {}, {}".format(shift, category) # if multiple shifts, build envelope if len(shifts) != 1: # shift with effect in up/down direction error_up = max([diff_up, diff_down, 0]) error_down = min([diff_up, diff_down, 0]) # add in quadrature if shift_idx == 0: errors_up.append(error_up**2) errors_down.append(error_down**2) else: errors_up[bin_idx - 1] += error_up**2 errors_down[bin_idx - 1] += error_down**2 else: errors_up.append(diff_up) errors_down.append(-diff_down) # is subtracted later # multiple shifts have been added quadratically, take square root if len(shifts) != 1: errors_up = np.sqrt(errors_up) errors_down = np.sqrt(errors_down) # build shifted histograms combined_hist_up = nominal_hist.Clone() combined_hist_down = nominal_hist.Clone() for bin_idx in range(1, nominal_hist.GetNbinsX() + 1): combined_hist_up.SetBinContent( bin_idx, combined_hist_up.GetBinContent(bin_idx) + errors_up[bin_idx - 1]) combined_hist_down.SetBinContent( bin_idx, combined_hist_down.GetBinContent(bin_idx) - errors_down[bin_idx - 1]) if self.norm_to_nominal: combined_hist_up.Divide(nominal_hist) combined_hist_down.Divide(nominal_hist) plot = ROOTPlot(hist_name, hist_name) plot.create_pads() plot.cd(0, 0) plot.draw({"nominal": nominal_hist}, line_color=1) plot.draw({"up": combined_hist_up}, line_color=2) plot.draw({"down": combined_hist_down}, line_color=4) plot.save(os.path.join(local_tmp.path, "{}.pdf".format(hist_name))) ## ## Check scale factors, uncertainties, and fits ## for i, (pt_idx, eta_idx) in enumerate( itertools.product(range(n_pt_categories), range(n_eta_categories))): data_hist = root_file.Get("h_csv_Data_Pt{}_Eta{}".format( pt_idx, eta_idx)) if self.flavor == "hf": signal_base = "h_csv_MC_bjets" bg_base = "h_csv_MC_nonbjets" else: signal_base = "h_csv_MC_nonbjets" # actually lf bg_base = "h_csv_MC_bjets" # actually b + c signal_hist = root_file.Get("{}_Pt{}_Eta{}".format( signal_base, pt_idx, eta_idx)) bg_hist = root_file.Get("{}_Pt{}_Eta{}".format( bg_base, pt_idx, eta_idx)) # with self.output().localize("w") as tmp: with tarfile.open(tmp.path, "w:gz") as tar: for plot_file in os.listdir(local_tmp.path): tar.add(os.path.join(local_tmp.path, plot_file), arcname=plot_file)
def merge_parquet_task(task, inputs, output, local=False, cwd=None, force=True, writer_opts=None): """ This method is intended to be used by tasks that are supposed to merge parquet files, e.g. when inheriting from :py:class:`law.contrib.tasks.MergeCascade`. *inputs* should be a sequence of targets that represent the files to merge into *output*. When *local* is *False* and files need to be copied from remote first, *cwd* can be a set as the dowload directory. When empty, a temporary directory is used. The *task* itself is used to print and publish messages via its :py:meth:`law.Task.publish_message` and :py:meth:`law.Task.publish_step` methods. When *force* is *True*, any existing output file is overwritten. *writer_opts* is forwarded to :py:func:`merge_parquet_files` which is used internally for the actual merging. """ # ensure inputs are targets inputs = [ LocalFileTarget(inp) if isinstance(inp, six.string_types) else inp for inp in inputs ] # ensure output is a target if isinstance(output, six.string_types): output = LocalFileTarget(output) def merge(inputs, output): with task.publish_step("merging {} parquet files ...".format( len(inputs)), runtime=True): # clear the output if necessary if output.exists() and force: output.remove() if len(inputs) == 1: output.copy_from_local(inputs[0]) else: # merge merge_parquet_files([inp.path for inp in inputs], output.path, writer_opts=writer_opts) # print the size output_size = human_bytes(output.stat().st_size, fmt=True) task.publish_message(f"merged file size: {output_size}") if local: # everything is local, just merge merge(inputs, output) else: # when not local, we need to fetch files first into the cwd if not cwd: cwd = LocalDirectoryTarget(is_tmp=True) elif isinstance(cwd, str): cwd = LocalDirectoryTarget(cwd) cwd.touch() # fetch with task.publish_step("fetching inputs ...", runtime=True): def fetch(inp): local_inp = cwd.child(inp.unique_basename, type="f") inp.copy_to_local(local_inp, cache=False) return local_inp def callback(i): task.publish_message("fetch file {} / {}".format( i + 1, len(inputs))) local_inputs = map_verbose(fetch, inputs, every=5, callback=callback) # merge into a localized output with output.localize("w", cache=False) as local_output: merge(local_inputs, local_output)
def hadd_task(task, inputs, output, cwd=None, local=False, force=True): """ This method is intended to be used by tasks that are supposed to merge root files, e.g. when inheriting from :py:class:`law.contrib.tasks.MergeCascade`. *inputs* should be a sequence of local targets that represent the files to merge into *output*. *cwd* is the working directory in which hadd is invoked. When empty, a temporary directory is used. The *task* itself is used to print and publish messages via its :py:meth:`law.Task.publish_message` and :py:meth:`law.Task.publish_step` methods. When *local* is *True*, the input and output targets are assumed to be local and the merging is based on their local paths. Otherwise, the targets are fetched first and the output target is localized. When *force* is *True*, any existing output file is overwritten (by adding the ``-f`` flag to ``hadd``). """ # ensure inputs are targets inputs = [ LocalFileTarget(inp) if isinstance(inp, six.string_types) else inp for inp in inputs ] # ensure output is a target if isinstance(output, six.string_types): output = LocalFileTarget(output) # default cwd if not cwd: cwd = LocalDirectoryTarget(is_tmp=True) elif isinstance(cwd, six.string_types): cwd = LocalDirectoryTarget(cwd) cwd.touch() # helper to create the hadd cmd def hadd_cmd(input_paths, output_path): cmd = ["hadd", "-n", "0"] if force: cmd.append("-f") cmd.extend(["-d", cwd.path]) cmd.append(output_path) cmd.extend(input_paths) return quote_cmd(cmd) if local: # when local, there is no need to download inputs input_paths = [inp.path for inp in inputs] with task.publish_step("merging ...", runtime=True): if len(inputs) == 1: output.copy_from_local(inputs[0]) else: # merge using hadd cmd = hadd_cmd(input_paths, output.path) code = interruptable_popen(cmd, shell=True, executable="/bin/bash")[0] if code != 0: raise Exception("hadd failed") task.publish_message("merged file size: {}".format(human_bytes( output.stat.st_size, fmt=True))) else: # when not local, we need to fetch files first into the cwd with task.publish_step("fetching inputs ...", runtime=True): def fetch(inp): inp.copy_to_local(cwd.child(inp.unique_basename, type="f"), cache=False) return inp.unique_basename def callback(i): task.publish_message("fetch file {} / {}".format(i + 1, len(inputs))) bases = map_verbose(fetch, inputs, every=5, callback=callback) # start merging into the localized output with output.localize("w", cache=False) as tmp_out: with task.publish_step("merging ...", runtime=True): if len(bases) == 1: tmp_out.path = cwd.child(bases[0]).path else: # merge using hadd cmd = hadd_cmd(bases, tmp_out.path) code = interruptable_popen(cmd, shell=True, executable="/bin/bash", cwd=cwd.path)[0] if code != 0: raise Exception("hadd failed") task.publish_message("merged file size: {}".format(human_bytes( tmp_out.stat.st_size, fmt=True)))
def htcondor_output_directory(self): # the directory where submission meta data should be stored return LocalDirectoryTarget(self.local_path())
def run(self): def add_hist(hist, new_hist, sign=1.): if hist is None: hist = new_hist.Clone() hist.Scale(sign) else: hist.Add(new_hist, sign) return hist import ROOT ROOT.PyConfig.IgnoreCommandLineOptions = True ROOT.gROOT.SetBatch() inp = self.input() outp = self.output() if self.normalize: scales = inp["scale"]["channel_scales"].load() local_tmp = LocalDirectoryTarget(is_tmp=True) local_tmp.touch() categories = [] for category, _, _ in self.config_inst.walk_categories(): if category.has_tag((self.category_tag, self.b_tagger), mode=all): categories.append(category) # create plot objects plot_dict = {} for category in categories: plot = ROOTPlot(category.name, category.name) plot.create_pads(n_pads_y=2, limits_y=[0., 0.3, 1.0], legend_loc="upper") plot_dict[category] = plot with inp["hists"].load("r") as input_file: for category in categories: data_hist = None mc_hists = defaultdict(lambda: defaultdict(lambda: None)) # shift -> key (process/flavor) for leaf_cat, _, children in category.walk_categories(): # we are only interested in leaves if children: continue flavor = leaf_cat.get_aux("flavor", None) channel = leaf_cat.get_aux("channel") region = leaf_cat.get_aux("region", None) category_dir = input_file.GetDirectory(leaf_cat.name) for process_key in category_dir.GetListOfKeys(): process = self.config_inst.get_process(process_key.GetName()) process_dir = category_dir.GetDirectory(process.name) # avoid double counting of inclusive and flavor-dependent histograms if flavor is not None: # Not needed in case region isn't flavor specific if process.is_data and flavor != "inclusive": continue elif process.is_mc and flavor == "inclusive": continue for shift in self.shifts: if process.is_data and shift != "nominal": continue for variable in self.variable: # create variable name from template aux = leaf_cat.aux.copy() aux["b_tag_var"] = self.config_inst.get_aux("btaggers")[self.b_tagger]["variable"] aux["b_tagger"] = self.b_tagger aux["shift"] = shift variable = variable.format(**aux) hist = process_dir.Get(variable) binning_type = "measurement" if self.rebin else None hist = self.rebin_hist(hist, region, binning_type=binning_type, truncate=self.truncate) add_to_data, sign = self.associate_hist(process=process, flavor=flavor, region=region) if add_to_data: if shift != "nominal": raise Exception("Cannot add shifted samples to data.") data_hist = add_hist(data_hist, hist, sign=sign) else: if self.normalize and region is not None: # apply "trigger" sfs as part of the normalization hist.Scale(scales[channel.name][region]) key = process if self.mc_split == "process" else flavor mc_hists[shift][key] = add_hist(mc_hists[shift][key], hist, sign=sign) if self.normalize: # normalize mc yield to data in this category mc_yield = sum(hist.Integral() for hist in mc_hists["nominal"].values()) data_yield = data_hist.Integral() norm_factor = data_yield / mc_yield for shift in self.shifts: for mc_hist in mc_hists[shift].values(): mc_hist.Scale(norm_factor) # get maximum value of hists/ stacks drawn to set axis ranges mc_hist_sum = mc_hists["nominal"].values()[0].Clone() for mc_hist in mc_hists["nominal"].values()[1:]: mc_hist_sum.Add(mc_hist) hist_maximum = max([mc_hist_sum.GetMaximum(), data_hist.GetMaximum()]) # get plot names mc_key = self.mc_key.format(**{"region": category.get_aux("region", None)}) data_key = self.data_key.format(**{"region": category.get_aux("region", None)}) plot = plot_dict[category] # data and mc histograms plot.cd(0, 1) if self.draw_stacked: plot.draw(mc_hists["nominal"], stacked=True, stack_maximum=1.5*hist_maximum, y_title="Entries") else: # fix axis range invis_hist = mc_hist_sum.Clone() if mc_hist_sum.GetMaximum() > data_hist.GetMaximum() else data_hist.Clone() invis_hist.Scale(1.5) plot.draw({"invis": invis_hist}, invis=True) plot.draw({mc_key: mc_hist_sum}, line_color=None) plot.draw({data_key: data_hist}) if self.draw_systematics: up_shifted_mc_hists = {} down_shifted_mc_hists = {} for shift in self.shifts: # combine processes/ flavors shifted_mc_hist_sum = mc_hists[shift].values()[0].Clone() for mc_hist in mc_hists[shift].values()[1:]: shifted_mc_hist_sum.Add(mc_hist) if shift.endswith("_down"): down_shifted_mc_hists[shift[:-5]] = shifted_mc_hist_sum.Clone() elif shift.endswith("_up"): up_shifted_mc_hists[shift[:-3]] = shifted_mc_hist_sum.Clone() envelope = build_hist_envelope(mc_hist_sum, up_shifted_mc_hists, down_shifted_mc_hists, envelope_as_errors=True) plot.draw_as_graph(envelope, options="2", hatched=True) # add category information to plot pt_range = category.get_aux("pt", None) eta_range = category.get_aux("eta", None) if pt_range is not None and eta_range is not None: if not np.isinf(pt_range[1]): text = r"#splitline{%d < p_{T} < %d}{%.1f < |#eta| < %.1f}" % \ (pt_range[0], pt_range[1], eta_range[0], eta_range[1]) else: text = r"#splitline{p_{T} > %d}{%.1f < |#eta| < %.1f}" % \ (pt_range[0], eta_range[0], eta_range[1]) plot.draw_text(text, size=0.05, xpos=0.505, ypos=0.5, align=11) # ratio of data to mc below the main plot plot.cd(0, 0) # ratio histograms # mc error band ratio_mcerr_hist = mc_hist_sum.Clone() # divide without error propagation self.divide_hists(ratio_mcerr_hist, mc_hist_sum) # ratio ratio_hist = data_hist.Clone() self.divide_hists(ratio_hist, mc_hist_sum) y_axis = ratio_hist.GetYaxis() y_axis.SetRangeUser(0.5, 1.5) y_axis.SetTitle("data/MC") y_axis.SetTitleSize(y_axis.GetTitleSize() * plot.open_pad.scale_factor) y_axis.SetLabelSize(y_axis.GetLabelSize() * plot.open_pad.scale_factor) y_axis.SetNdivisions(505) y_axis.SetTitleOffset(0.65) x_axis = ratio_hist.GetXaxis() if self.x_title: aux = category.aux.copy() aux["b_tag_var"] = self.config_inst.get_aux("btaggers")[self.b_tagger]["label"] x_axis.SetTitle(self.x_title.format(**aux)) x_axis.SetTitleSize(x_axis.GetTitleSize() * plot.open_pad.scale_factor) x_axis.SetLabelSize(x_axis.GetLabelSize() * plot.open_pad.scale_factor) plot.draw({"invis": ratio_hist}, invis=True) plot.draw_as_graph(ratio_mcerr_hist, options="2") plot.draw({"data/mc": ratio_hist}) if self.draw_systematics: # build envelope of ratio to nominal hist for hist in up_shifted_mc_hists.values(): hist.Divide(mc_hist_sum) for hist in down_shifted_mc_hists.values(): hist.Divide(mc_hist_sum) scaled_envelope = build_hist_envelope(ratio_mcerr_hist, up_shifted_mc_hists, down_shifted_mc_hists, envelope_as_errors=True) plot.draw_as_graph(scaled_envelope, options="2", hatched=True) for category, plot in plot_dict.items(): plot_name = self.get_plot_name(category.name, self.variable, self.b_tagger, self.iteration) plot.save(os.path.join(local_tmp.path, plot_name), draw_legend=(False, True), log_y=self.logarithmic, lumi=self.config_inst.get_aux("lumi").values()[0]/1000.) del plot with outp.localize("w") as tmp: with tarfile.open(tmp.path, "w:gz") as tar: for plot_file in os.listdir(local_tmp.path): tar.add(os.path.join(local_tmp.path, plot_file), arcname=plot_file)
def output(self): return LocalDirectoryTarget(self.path)