tracking-parametrisation-tuner/parameterisations/parameterise_magnet_kink_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_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"]

    # array = array[(array["z_mag_xEndT"] < 5500) & (array["z_mag_xEndT"] > 5000)]

    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