tracking-parametrisation-tuner/parameterisations/parameterise_field_integral.py

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