tracking-parametrisation-tuner/parameterisations/parameterise_track_model_electron.py

# 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_track_model(
    input_file: str = "data/tracking_losses_ntuple_B_default_selected.root",
    tree_name: str = "Selected",
) -> Path:
    """Function that calculates the parameterisations to estimate track model coefficients.

    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 files containing the found parameters.
    """
    input_tree = uproot.open({input_file: tree_name})
    # this is an event list of dictionaries containing awkward arrays
    array = input_tree.arrays()

    array["yStraightOut"] = array["ideal_state_770_y"] + array["ideal_state_770_ty"] * (
        array["ideal_state_10000_z"] - array["ideal_state_770_z"]
    )
    array["yDiffOut"] = array["ideal_state_10000_y"] - array["yStraightOut"]

    array = array[(array["yDiffOut"] < 100) & (array["yDiffOut"] > -100)]

    array["yStraightEndT"] = array["ideal_state_770_y"] + array[
        "ideal_state_770_ty"
    ] * (9410.0 - array["ideal_state_770_z"])
    array["yDiffEndT"] = array["ideal_state_9410_y"] - array["yStraightEndT"]

    array["dSlope_xEndT"] = array["ideal_state_9410_tx"] - array["ideal_state_770_tx"]
    array["dSlope_yEndT"] = array["ideal_state_9410_ty"] - array["ideal_state_770_ty"]
    array["dSlope_xEndT_abs"] = abs(array["dSlope_xEndT"])
    array["dSlope_yEndT_abs"] = abs(array["dSlope_yEndT"])

    model_y_match, poly_features_y_match = fit_linear_regression_model(
        array,
        target_feat="yDiffOut",
        features=[
            "ideal_state_770_ty",
            "dSlope_xEndT",
            "dSlope_yEndT",
        ],
        keep=[
            "dSlope_yEndT",
            "dSlope_xEndT dSlope_yEndT",
            "ideal_state_770_ty dSlope_xEndT^2",
            "ideal_state_770_ty dSlope_yEndT^2",
        ],
        degree=3,
        fit_intercept=False,
    )
    keep_y_match_precise = [
        "dSlope_yEndT",
        "ideal_state_770_ty dSlope_xEndT_abs",
        "ideal_state_770_ty dSlope_yEndT_abs",
        "ideal_state_770_ty dSlope_yEndT^2",
        "ideal_state_770_ty dSlope_xEndT^2",
        "ideal_state_770_ty ideal_state_770_tx dSlope_xEndT",
        "ideal_state_770_tx^2 dSlope_yEndT",
        "ideal_state_770_ty ideal_state_770_tx^2 dSlope_xEndT_abs",
        "ideal_state_770_ty^3 ideal_state_770_tx dSlope_xEndT",
    ]
    model_y_match_precise, poly_features_y_match_precise = fit_linear_regression_model(
        array,
        "yDiffEndT",
        [
            "ideal_state_770_ty",
            "ideal_state_770_tx",
            "dSlope_xEndT",
            "dSlope_yEndT",
            "dSlope_xEndT_abs",
            "dSlope_yEndT_abs",
        ],
        keep=keep_y_match_precise,
        degree=5,
    )

    cpp_y_match = parse_regression_coef_to_array(
        model_y_match,
        poly_features_y_match,
        array_name="bendYParamsMatchElectron",
    )
    cpp_y_match_precise = parse_regression_coef_to_array(
        model_y_match_precise,
        poly_features_y_match_precise,
        array_name="bendYParamsElectron",
    )

    outpath = Path("parameterisations/result/track_model_params_electron.hpp")
    outpath.parent.mkdir(parents=True, exist_ok=True)
    with open(outpath, "w") as result:
        result.writelines(
            cpp_y_match + cpp_y_match_precise,
        )
    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_track_model(**args_dict)

    try:
        import subprocess

        # run clang-format for nicer looking result
        subprocess.run(
            [
                "clang-format",
                "-i",
                f"{outfile}",
            ],
            check=True,
        )
    except:
        pass