You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

94 lines
3.4 KiB

import uproot
import awkward as ak
from pathlib import Path
from scipy.optimize import curve_fit
import numpy as np
import pandas as pd
from math import ceil
import argparse
def fastSigmoid(x, p0, p1, p2):
return p0 + p1 * x / (1 + abs(p2 * x))
def parameterise_hough_histogram(
input_file: str = "data/param_data_selected_all_p.root",
tree_name: str = "Selected",
n_bins_start: int = 900,
hist_range: tuple[float, float] = (-3000.0, 3000.0),
first_bin_center: float = 2.5,
) -> Path:
"""Function to parameterise the binning of the Hough histogram using the occupancy on the reference plane.
Args:
input_file (str, optional): Defaults to "data/param_data_selected_all_p.root".
tree_name (str, optional): Defaults to "Selected".
n_bins_start (int, optional): Starting (minimal) number of bins in histogram. Defaults to 900.
hist_range (tuple[float, float], optional): Range in mm the histogram covers. Defaults to (-3000.0, 3000.0).
first_bin_center (float, optional): Calculated bin center at lower range. Defaults to 2.5.
Returns:
Path: Path to cpp code file.
"""
input_tree = uproot.open({input_file: tree_name})
# this is an event list of dictionaries containing awkward arrays
array = input_tree.arrays()
array = array[[field for field in ak.fields(array) if "scifi_hit" not in field]]
df = ak.to_pandas(array)
selection = (df["x_ref"] > hist_range[0]) & (df["x_ref"] < hist_range[1])
data = df.loc[selection, "x_ref"]
_, equal_bins = pd.qcut(data, q=n_bins_start, retbins=True)
bin_numbering = np.arange(0, n_bins_start + 1)
equalbins_center = equal_bins[int(n_bins_start / 10) : int(9 * n_bins_start / 10)]
bin_numbering_center = bin_numbering[
int(n_bins_start / 10) : int(9 * n_bins_start / 10)
]
func = fastSigmoid
popt, _ = curve_fit(func, xdata=equalbins_center, ydata=bin_numbering_center)
print("Parameterisation for central occupancy:")
print("fastSigmoid(x,", *popt, ")")
print("Scan shift to match first bin center ...")
shift = 0.0
while func(hist_range[0], popt[0] + shift, *popt[1:]) < first_bin_center:
shift += 0.1
popt[0] += shift - 0.1
popt[2] = abs(popt[2])
print("shifted: fastSigmoid(x,", *popt, ")")
n_bins_final = ceil(func(hist_range[1], *popt))
print(
"Final number of bins:",
n_bins_final,
"including offset of",
int(first_bin_center),
)
comment = f"// p[0] + p[1] * x / (1 + p[2] * abs(x)) for nBins = {n_bins_final}\n"
cpp = (
"constexpr auto p = std::array{"
+ ", ".join([str(p) + "f" for p in popt])
+ "};"
)
outpath = Path("parameterisations/result/hough_histogram_binning_params.hpp")
outpath.parent.mkdir(parents=True, exist_ok=True)
with open(outpath, "w") as result:
result.writelines([comment, cpp])
return outpath
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--input-file",
type=str,
help="Path to the input file",
required=False,
)
parser.add_argument(
"--tree-name",
type=str,
help="Path to the input file",
required=False,
)
args = parser.parse_args()
args_dict = {arg: val for arg, val in vars(args).items() if val is not None}
outfile = parameterise_hough_histogram(**args_dict)