Add data preparation script

prepare.py

+
+# Copyright 2022, Frank Sauerburger
+
+"""Module to prepare the ROOT file for the PyHEP conference"""
+
+import argparse
+import array
+import requests
+import ROOT
+
+DATA_URL = "https://gitlab.cern.ch/fsauerbu/uhepp/-/raw/master/docs/toyhisto.json"
+
+def download():
+    """Download the example data file"""
+    toydata = requests.get(DATA_URL).json()
+    return toydata
+
+
+def create_th1f(bin_edges, base, stat):
+    """Create a histogram from the central values and stat uncertainties"""
+    n_bins = len(bin_edges) - 1
+    bins = array.array('f', bin_edges)
+    histo = ROOT.TH1F("", "", n_bins, bins)
+    for i, (base_val, stat_val) in enumerate(zip(base, stat)):
+        histo.SetBinContent(i + 1, base_val)
+        histo.SetBinError(i + 1, stat_val)
+
+    return histo
+
+
+def create_file(toy_data, filename):
+    """Create toy ROOT file"""
+    root_file = ROOT.TFile.Open(filename, "RECREATE")
+    
+    bin_edges = toy_data["bin_edges"]
+    sig = create_th1f(bin_edges, toy_data["sig"], toy_data["sig_stat"])
+    sig.SetName("signal")
+
+    data = create_th1f(bin_edges, toy_data["data"], toy_data["data_stat"])
+    data.SetName("data")
+
+    bkg = create_th1f(bin_edges, toy_data["bkg"], toy_data["bkg_stat"])
+    bkg.SetName("bkg")
+
+    root_file.Write()
+    root_file.Close()
+
+
+def default_parser():
+    """Return default argument parser"""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("output", metavar="OUTFILE", default="data.root",
+                        help="Output root file")
+    return parser
+
+
+def main(args):
+    """Main CLI entrypoint"""
+    toy_data = download()
+    create_file(toy_data, args.output)
+
+
+if __name__ == "__main__":
+    parser = default_parser()
+    args = parser.parse_args()
+    main(args)