read_file.py

import ROOT
import numpy as np
import time


def arg_parse():
    import argparse

    parser = argparse.ArgumentParser(description='Process some integers.')
    parser.add_argument('-i', type=str, required=True, help='Input file string')
    parser.add_argument('-d', type=str, required=True, help='Output directory')
    parser.add_argument('-isotope', type=str, required=True, help='Isotope')
    args = parser.parse_args()
    return args


def main():
    args = arg_parse()
    input_file = args.i
    output_directory = args.d
    isotope = args.isotope

    sim_type = input_file.split("/")[-1].split("_")[0].strip()
    thickness = int(input_file.split("/")[-1].split("_")[-1].split(".")[0].strip())

    output_filename = output_directory + "/{}_{}um_output.root".format(sim_type, thickness)

    try:
        file = open(input_file, "r")
        print(">>> Reading file: {}".format(input_file))
        file.close()
    except (OSError, IOError):
        print(">>> Failed to open file: {}".format(input_file))
        return 1
    file = ROOT.TFile(input_file, "READ")
    tree = file.Get("SimData")

    try:
        output_file = ROOT.TFile(output_filename, "RECREATE")
    except (OSError, IOError):
        print(">>> Failed to open file: {}".format(output_filename))
        return 1

    # Define a series of TH1D histograms
    if isotope == 'Nd150':
        max_energy = 3.4  # MeV
        min_energy = 0
    elif isotope == 'Se82':
        max_energy = 3  # MeV
        min_energy = 0
    else:
        max_energy = 3  # MeV
        min_energy = 0
    energy_bin_width = 0.005
    n_bins = int((max_energy - min_energy)/energy_bin_width)
    e_bins = np.array([(i + 0.5)*energy_bin_width for i in range(n_bins)])

    sim_energy_vertex_hist = ROOT.TH1D("sim_energy_vertex_hist", "sim_energy_vertex_hist",
                                       n_bins, min_energy, max_energy)
    sim_energy_foil_hist = ROOT.TH1D("sim_energy_foil_hist", "sim_energy_foil_hist",
                                     n_bins, min_energy, max_energy)
    summed_sim_energy_vertex_hist = ROOT.TH1D("summed_sim_energy_vertex_hist", "summed_sim_energy_vertex_hist",
                                              n_bins, min_energy, max_energy)
    summed_sim_energy_foil_hist = ROOT.TH1D("summed_sim_energy_foil_hist", "summed_sim_energy_foil_hist",
                                            n_bins, min_energy, max_energy)
    summed_sim_energy_foil_hist_same = ROOT.TH1D("summed_sim_energy_foil_hist_same", "summed_sim_energy_foil_hist_same",
                                                 n_bins, min_energy, max_energy)
    summed_sim_energy_foil_hist_opposite = ROOT.TH1D("summed_sim_energy_foil_hist_opposite",
                                                     "summed_sim_energy_foil_hist_opposite",
                                                     n_bins, min_energy, max_energy)
    summed_sim_energy_foil_hist_rejected = ROOT.TH1D("summed_sim_energy_foil_hist_rejected",
                                                     "summed_sim_energy_foil_hist_rejected",
                                                     n_bins, min_energy, max_energy)
    summed_sim_energy_foil_hist_total = ROOT.TH1D("summed_sim_energy_foil_hist_total",
                                                  "summed_sim_energy_foil_hist_total",
                                                  n_bins, min_energy, max_energy)
    displacement_hist = ROOT.TH1D("displacement_hist", "displacement_hist",
                                  100, -thickness/2000, thickness/2000)
    efficiency_hist = ROOT.TH1D("efficiency_hist", "efficiency_hist",
                                n_bins, min_energy, max_energy)
    sim_energy_foil_hists = []
    energy_loss_hists = []
    for i in range(n_bins):
        the_max = (i+1)*energy_bin_width
        sim_energy_foil_hists.append(ROOT.TH1D("sim_energy_foil_hist_{}keV_bin".format(int(round(the_max*1000, 0))),
                                               "sim_energy_foil_hist_{}keV_bin".format(int(round(the_max*1000, 0))),
                                               int(the_max/energy_bin_width), 0, the_max))
        energy_loss_hists.append(ROOT.TH1D("energy_loss_hist_{}keV_bin".format(int(round(the_max*1000, 0))),
                                           "energy_loss_hist_{}keV_bin".format(int(round(the_max*1000, 0))),
                                           int(the_max/energy_bin_width), 0, the_max))

    i_event = 0
    previous_time = time.time()
    for event in tree:
        if (i_event + 1) % 100000 == 0:
            t = time.time()
            print(">>> Progress: {}/{} {:.2f}% {:.2f}s".format(i_event, tree.GetEntries(),
                                                               100 * (i_event + 1) / tree.GetEntries(),
                                                               t - previous_time))
            previous_time = t

        escape_event = True
        sides_event = True
        summed_vertex_energy = 0
        summed_foil_energy = 0

        for index in range(len(event.og_energy)):
            vertex_energy = event.og_energy[index]
            energy_loss = event.energy_loss[index]
            foil_energy = vertex_energy - energy_loss
            if foil_energy == 0.0:
                escape_event = False
                eff_energy_bin = int(vertex_energy/energy_bin_width)
                val = efficiency_hist.GetBinContent(eff_energy_bin)
                efficiency_hist.SetBinContent(eff_energy_bin, val + 1)

            displacement = event.displacement[index]

            sim_energy_vertex_hist.Fill(vertex_energy)
            sim_energy_foil_hist.Fill(foil_energy)
            displacement_hist.Fill(displacement)

            i_bin = int(vertex_energy/energy_bin_width)
            sim_energy_foil_hists[i_bin].Fill(foil_energy)
            energy_loss_hists[i_bin].Fill(energy_loss)

            summed_vertex_energy += vertex_energy
            summed_foil_energy += foil_energy

        index0 = np.where(np.array(list(event.step_track_id)) == event.track_ids[0])[0][-1]
        index1 = np.where(np.array(list(event.step_track_id)) == event.track_ids[1])[0][-1]
        last_x0 = list(event.step_stop_vertex_x)[index0]
        last_x1 = list(event.step_stop_vertex_x)[index1]

        t = thickness/1000
        if abs(last_x0) < t / 2 and abs(last_x1) < t / 2:
            sides_event = False

        if (last_x0 < 0 and last_x1 < 0) or (last_x0 > 0 and last_x1 > 0):
            case = 0
        else:
            case = 1

        summed_sim_energy_vertex_hist.Fill(summed_vertex_energy)
        if escape_event and sides_event:
            summed_sim_energy_foil_hist.Fill(summed_foil_energy)
            if case == 0:
                summed_sim_energy_foil_hist_same.Fill(summed_foil_energy)
            if case == 1:
                summed_sim_energy_foil_hist_opposite.Fill(summed_foil_energy)
        else:
            summed_sim_energy_foil_hist_rejected.Fill(summed_foil_energy)
        summed_sim_energy_foil_hist_total.Fill(summed_foil_energy)

        i_event += 1

    for i_bin in range(efficiency_hist.GetNbinsX()):
        val = efficiency_hist.GetBinContent(i_bin)
        norm = sim_energy_vertex_hist.GetBinContent(i_bin)
        if norm == 0:
            efficiency_hist.SetBinContent(i_bin, 0)
        else:
            efficiency_hist.SetBinContent(i_bin, val/norm)

    print(">>> Writing to file: {} events: {}".format(output_filename, tree.GetEntries()))

    file.Close()
    output_file.cd()
    output_file.Write()
    output_file.Close()
    print(">>> Finished")

    return 0


if __name__ == "__main__":
    err_code = main()
    exit(err_code)