109 lines
3.3 KiB
Python
109 lines
3.3 KiB
Python
|
# flake8: noqa
|
||
|
|
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_magnet_kink(
|
||
|
|
input_file: str = "data/tracking_losses_ntuple_B_default_selected.root",
|
||
|
|
tree_name: str = "Selected",
|
||
|
|
) -> Path:
|
||
|
|
"""Function that calculates parameters for estimating the magnet kink z position.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
input_file (str, optional): Defaults to "data/param_data_selected.root".
|
||
|
|
tree_name (str, optional): Defaults to "Selected".
|
||
|
|
per_layer (bool, optional): If true also calculates parameters per SciFi layer. Defaults to False.
|
||
|
|
|
||
|
|
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 = allarray[
|
||
|
|
# (allarray.ideal_state_770_x == allarray.ideal_state_770_x)
|
||
|
|
# & (allarray.ideal_state_9410_x == allarray.ideal_state_9410_x)
|
||
|
|
# & (allarray.ideal_state_10000_x == allarray.ideal_state_10000_x)
|
||
|
|
# ]
|
||
|
|
array["dSlope_xEndT"] = array["ideal_state_9410_tx"] - array["ideal_state_770_tx"]
|
||
|
|
array["dSlope_xEndT_abs"] = abs(array["dSlope_xEndT"])
|
||
|
|
array["x_EndT_abs"] = abs(array["ideal_state_9410_x"])
|
||
|
|
# the magnet kink position is the point of intersection of the track tagents
|
||
|
|
array["z_mag_xEndT"] = (
|
||
|
|
array["ideal_state_770_x"]
|
||
|
|
- array["ideal_state_9410_x"]
|
||
|
|
- array["ideal_state_770_tx"] * array["ideal_state_770_z"]
|
||
|
|
+ array["ideal_state_9410_tx"] * array["ideal_state_9410_z"]
|
||
|
|
) / array["dSlope_xEndT"]
|
||
|
|
|
||
|
|
model_endt, poly_features_endt = fit_linear_regression_model(
|
||
|
|
array,
|
||
|
|
target_feat="z_mag_xEndT",
|
||
|
|
features=[
|
||
|
|
"ideal_state_770_tx",
|
||
|
|
"dSlope_xEndT",
|
||
|
|
"dSlope_xEndT_abs",
|
||
|
|
"x_EndT_abs",
|
||
|
|
],
|
||
|
|
keep=[
|
||
|
|
"ideal_state_770_tx^2",
|
||
|
|
"dSlope_xEndT^2",
|
||
|
|
"dSlope_xEndT_abs",
|
||
|
|
"x_EndT_abs",
|
||
|
|
],
|
||
|
|
degree=2,
|
||
|
|
fit_intercept=True,
|
||
|
|
)
|
||
|
|
cpp_ref = parse_regression_coef_to_array(
|
||
|
|
model_endt,
|
||
|
|
poly_features_endt,
|
||
|
|
"zMagnetParamsEndT",
|
||
|
|
)
|
||
|
|
|
||
|
|
outpath = Path("parameterisations/result/z_mag_kink_params_electron.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_magnet_kink(**args_dict)
|
||
|
|
|
||
|
|
try:
|
||
|
|
import subprocess
|
||
|
|
|
||
|
|
# run clang-format for nicer looking result
|
||
|
|
subprocess.run(
|
||
|
|
[
|
||
|
|
"clang-format",
|
||
|
|
"-i",
|
||
|
|
f"{outfile}",
|
||
|
|
],
|
||
|
|
check=True,
|
||
|
|
)
|
||
|
|
except:
|
||
|
|
pass
|