95 lines
3.4 KiB
Python
95 lines
3.4 KiB
Python
|
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)
|