tracking-parametrisation-tuner/parameterisations/parameterise_search_window.py

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