Working with ROOT files in law.

law_root_examples.py

# coding: utf-8

# This script showcases a few examples on how to simplify working with ROOT objects within law.
#
# Scope:
#   1. Usage of the low-level GuardedTFile
#   2. Reading and writing ROOT files via target formatters
#   3. Merging ROOT files with hadd_task

import law
import luigi


# first, load the root contrib package
# this makes objects exposed in https://github.com/riga/law/blob/master/law/contrib/root/__init__.py
# available under law.root
law.contrib.load("root")

# import ROOT the least intrusive way possible
ROOT = law.root.import_ROOT()


# 1. Usage of the low-level GuardedTFile

# write a histogram into a file
with law.root.GuardedTFile("law_test_file1.root", "RECREATE") as f:
    h = ROOT.TH1F("h", "", 10, 0., 10.)
    h.Fill(5)
    h.Write()

# read a histogram from a file
with law.root.GuardedTFile("law_test_file1.root", "READ") as f:
    h = f.Get("h")
    print("integral: {}".format(h.Integral()))


# 2. Reading and writing ROOT files via target formatters (implemented in load and dump)

# write a histogram into a file described by a local file target
# (note that the same api applies for remote targets)
t = law.LocalFileTarget("law_test_file2.root")
with t.dump("RECREATE", formatter="root") as f:  # dump is required by mode "RECREATE"
    h = ROOT.TH1F("h", "", 10, 0., 10.)
    h.Fill(7)
    h.Write()

# read a histogram from a file described by a local file target
# (note that the same api applies for remote targets)
with t.load("READ", formatter="root") as f:  # load is required by mode "READ"
    h = f.Get("h")
    print("integral: {}".format(h.Integral()))

# use the uproot formatter now
with t.load(formatter="uproot") as f:
    f["h"].show()


# 3. Merging ROOT files with hadd_task

class MergingTask(law.Task):

    n_inputs = luigi.IntParameter()

    def output(self):
        return law.LocalFileTarget("law_test_file_merged.root")

    @law.decorator.safe_output  # deletes the output when an exception is raised inside run
    def run(self):
        inputs = ["law_test_file{}.root".format(i + 1) for i in range(self.n_inputs)]
        output = self.output()

        # run the task version of hadd that includes the handling of remote files, temporary
        # directories, and more
        law.root.hadd_task(self, inputs, output, local=True)


# invoke the task which would normally be done via the command line, delete the output first
task = MergingTask(n_inputs=2)
task.output().remove()
task.law_run()

# show the merged histogram using the uproot formatter as done above
with task.output().load(formatter="uproot") as f:
    f["h"].show()