102 lines
3.0 KiB
Python
102 lines
3.0 KiB
Python
|
from parameterisations.utils.parse_regression_coef_to_array import (
|
||
|
|
parse_regression_coef_to_array,
|
||
|
|
)
|
||
|
|
from parameterisations.utils.fit_linear_regression_model import (
|
||
|
|
fit_linear_regression_model,
|
||
|
|
)
|
||
|
|
import uproot
|
||
|
|
import argparse
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
def parameterise_search_window(
|
||
|
|
input_file: str = "data/param_data_selected_all_p.root",
|
||
|
|
tree_name: str = "Selected",
|
||
|
|
) -> Path:
|
||
|
|
"""Function that calculates parameters for estimating the hit search window border.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
input_file (str, optional): Defaults to "data/param_data_selected.root".
|
||
|
|
tree_name (str, optional): Defaults to "Selected".
|
||
|
|
|
||
|
|
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["x_straight_diff_ref"] = (
|
||
|
|
array["x"] + array["tx"] * (array["z_ref"] - array["z"]) - array["x_ref"]
|
||
|
|
)
|
||
|
|
array["x_straight_diff_ref_abs"] = abs(array["x_straight_diff_ref"])
|
||
|
|
array["inv_p_gev"] = 1000.0 / array["p"]
|
||
|
|
array["pol_qop_gev"] = array["signed_rel_current"] * array["qop"] * 1000.0
|
||
|
|
|
||
|
|
# now fit model for the reference plane
|
||
|
|
model_ref, poly_features_ref = fit_linear_regression_model(
|
||
|
|
array,
|
||
|
|
target_feat="x_straight_diff_ref_abs",
|
||
|
|
features=["ty", "tx", "inv_p_gev", "pol_qop_gev"],
|
||
|
|
keep=[
|
||
|
|
"inv_p_gev",
|
||
|
|
"ty^2 inv_p_gev",
|
||
|
|
"tx^2 inv_p_gev",
|
||
|
|
"tx inv_p_gev pol_qop_gev",
|
||
|
|
"inv_p_gev^3",
|
||
|
|
"tx^3 pol_qop_gev",
|
||
|
|
"tx^2 inv_p_gev^2",
|
||
|
|
"tx inv_p_gev^2 pol_qop_gev",
|
||
|
|
"inv_p_gev^4",
|
||
|
|
"ty^2 tx^2 inv_p_gev",
|
||
|
|
"ty^2 tx inv_p_gev pol_qop_gev",
|
||
|
|
"ty^2 inv_p_gev^3",
|
||
|
|
"tx^4 inv_p_gev",
|
||
|
|
],
|
||
|
|
degree=5,
|
||
|
|
fit_intercept=False,
|
||
|
|
)
|
||
|
|
cpp_ref = parse_regression_coef_to_array(
|
||
|
|
model_ref,
|
||
|
|
poly_features_ref,
|
||
|
|
"momentumWindowParamsRef",
|
||
|
|
)
|
||
|
|
outpath = Path("parameterisations/result/search_window_params.hpp")
|
||
|
|
outpath.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
with open(outpath, "w") as result:
|
||
|
|
result.writelines(cpp_ref)
|
||
|
|
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_search_window(**args_dict)
|
||
|
|
|
||
|
|
try:
|
||
|
|
import subprocess
|
||
|
|
|
||
|
|
# run clang-format for nicer looking result
|
||
|
|
subprocess.run(
|
||
|
|
[
|
||
|
|
"clang-format",
|
||
|
|
"-i",
|
||
|
|
f"{outfile}",
|
||
|
|
],
|
||
|
|
check=True,
|
||
|
|
)
|
||
|
|
except:
|
||
|
|
pass
|