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.
99 lines
2.7 KiB
99 lines
2.7 KiB
from parameterisations.utils.fit_linear_regression_model import (
|
|
fit_linear_regression_model,
|
|
)
|
|
from parameterisations.utils.parse_regression_coef_to_array import (
|
|
parse_regression_coef_to_array,
|
|
)
|
|
import uproot
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
|
|
def parameterise_field_integral(
|
|
input_file: str = "data/param_data_selected_all_p.root",
|
|
tree_name: str = "Selected",
|
|
) -> Path:
|
|
"""Function to estimate parameters describing the magnetic field integral.
|
|
|
|
Args:
|
|
input_file (str, optional): Defaults to "data/param_data_selected_all_p.root".
|
|
tree_name (str, optional): Defaults to "Selected".
|
|
|
|
Returns:
|
|
Path: Path to the parameters in cpp format.
|
|
"""
|
|
input_tree = uproot.open({input_file: tree_name})
|
|
# this is an event list of dictionaries containing awkward arrays
|
|
array = input_tree.arrays()
|
|
array["dSlope_fringe"] = array["tx_ref"] - array["tx"]
|
|
array["poqmag_gev"] = 1.0 / (array["signed_rel_current"] * array["qop"] * 1000.0)
|
|
array["B_integral"] = array["poqmag_gev"] * array["dSlope_fringe"]
|
|
|
|
model_ref, poly_features_ref = fit_linear_regression_model(
|
|
array,
|
|
target_feat="B_integral",
|
|
features=[
|
|
"ty",
|
|
"tx",
|
|
"tx_ref",
|
|
],
|
|
keep=[
|
|
"ty^2",
|
|
"tx^2",
|
|
"tx tx_ref",
|
|
"tx_ref^2",
|
|
"ty^2 tx tx_ref",
|
|
"ty^2 tx^2",
|
|
"ty^2 tx_ref^2",
|
|
"tx^4",
|
|
"ty^4",
|
|
"tx_ref^4",
|
|
"tx^3 tx_ref",
|
|
],
|
|
degree=5,
|
|
fit_intercept=True,
|
|
)
|
|
cpp_ref = parse_regression_coef_to_array(
|
|
model_ref,
|
|
poly_features_ref,
|
|
"fieldIntegralParamsRef",
|
|
)
|
|
outpath = Path("parameterisations/result/field_integral_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_field_integral(**args_dict)
|
|
|
|
try:
|
|
import subprocess
|
|
|
|
# run clang-format for nicer looking result
|
|
subprocess.run(
|
|
[
|
|
"clang-format",
|
|
"-i",
|
|
f"{outfile}",
|
|
],
|
|
check=True,
|
|
)
|
|
except:
|
|
pass
|