tracking-parametrisation-tuner/parameterisations/parameterise_magnet_kink.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
import awkward as ak
from pathlib import Path


def parameterise_magnet_kink(
    input_file: str = "data/param_data_selected.root",
    tree_name: str = "Selected",
    per_layer=False,
) -> 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["dSlope_fringe"] = array["tx_ref"] - array["tx"]
    # the magnet kink position is the point of intersection of the track tagents
    array["z_mag_x_fringe"] = (
        array["x"]
        - array["x_ref"]
        - array["tx"] * array["z"]
        + array["tx_ref"] * array["z_ref"]
    ) / array["dSlope_fringe"]
    array["dSlope_xEndT"] = array["tx_l11"] - array["tx"]
    array["dSlope_xEndT_abs"] = abs(array["dSlope_xEndT"])
    array["x_EndT_abs"] = abs(
        array["x_l11"] + array["tx_l11"] * (9410.0 - array["z_l11"]),
    )
    # the magnet kink position is the point of intersection of the track tagents
    array["z_mag_xEndT"] = (
        array["x"]
        - array["x_l11"]
        - array["tx"] * array["z"]
        + array["tx_l11"] * array["z_l11"]
    ) / array["dSlope_xEndT"]

    if per_layer:
        layered_features = [f"x_diff_straight_l{layer}" for layer in range(12)]
        rows = []
        for i, feat in enumerate(layered_features):
            scale = 3000
            if "dSlope" not in feat:
                array[f"x_l{i}_rel"] = array[f"x_l{i}"] / scale
                array[f"x_diff_straight_l{i}"] = (
                    array[f"x_l{i}"]
                    - array["x"]
                    - array["tx"] * (array[f"z_l{i}"] - array["z"])
                )

            model, poly_features = fit_linear_regression_model(
                array,
                target_feat="z_mag_x_fringe",
                features=[
                    "tx",
                    "ty",
                    feat,
                ],
                keep=[
                    "tx^2",
                    f"tx x_diff_straight_l{i}",
                    "ty^2",
                    f"x_diff_straight_l{i}^2",
                ],
                degree=2,
                fit_intercept=True,
            )
            rows.append(
                "{"
                + str(model.intercept_)
                + "f,"
                + ",".join([str(coef) + "f" for coef in model.coef_ if coef != 0.0])
                + "}",
            )

        cpp_decl = parse_regression_coef_to_array(
            model,
            poly_features,
            "zMagnetParamsLayers",
            rows=rows,
        )
    # now fit model for the reference plane
    model_ref, poly_features_ref = fit_linear_regression_model(
        array,
        target_feat="z_mag_x_fringe",
        features=["tx", "ty", "dSlope_fringe"],
        keep=["tx^2", "tx dSlope_fringe", "ty^2", "dSlope_fringe^2"],
        degree=2,
        fit_intercept=True,
    )
    cpp_ref = parse_regression_coef_to_array(
        model_ref,
        poly_features_ref,
        "zMagnetParamsRef",
    )

    model_endt, poly_features_endt = fit_linear_regression_model(
        array,
        target_feat="z_mag_xEndT",
        features=["tx", "dSlope_xEndT", "dSlope_xEndT_abs", "x_EndT_abs"],
        keep=["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.hpp")
    outpath.parent.mkdir(parents=True, exist_ok=True)
    with open(outpath, "w") as result:
        result.writelines(cpp_decl + cpp_ref if per_layer else 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