tracking-parametrisation-tuner/main.py

# flake8: noqaq
import os
import subprocess
import argparse
from parameterisations.parameterise_magnet_kink_electron import parameterise_magnet_kink
from parameterisations.parameterise_track_model_electron import parameterise_track_model
from parameterisations.parameterise_search_window import parameterise_search_window
from parameterisations.parameterise_field_integral import parameterise_field_integral
from parameterisations.parameterise_hough_histogram import parameterise_hough_histogram
from parameterisations.utils.preselection_trackinglosses import preselection
from parameterisations.train_forward_ghost_mlps import (
    train_default_forward_ghost_mlp,
    train_veloUT_forward_ghost_mlp,
)
from parameterisations.train_matching_ghost_mlps import train_matching_ghost_mlp
from parameterisations.utils.parse_tmva_matrix_to_array import (
    parse_tmva_matrix_to_array,
)


parser = argparse.ArgumentParser()
parser.add_argument(
    "--field-params",
    action="store_true",
    help="Enables determination of magnetic field parameterisations.",
)
parser.add_argument(
    "--forward-weights",
    action="store_true",
    help="Enables determination of weights used by neural networks.",
)
parser.add_argument(
    "--matching-weights",
    action="store_true",
    # default=True,
    help="Enables determination of weights used by neural networks.",
)
parser.add_argument(
    "-p",
    "--prepare",
    action="store_true",
    help="Enables preparation of data for matching.",
)
parser.add_argument(
    "--prepare-params-data",
    action="store_true",
    help="Enables preparation of data for magnetic field parameterisations.",
)
parser.add_argument(
    "--prepare-weights-data",
    action="store_true",
    help="Enables preparation of data for NN weight determination.",
)
args = parser.parse_args()

selected = "data/tracking_losses_ntuple_B_BJpsi_def_selected.root"
if args.prepare_params_data:
    selection = "isElectron == 1 && pt > 10 && p > 1500 && p < 100000 && !std::isnan(ideal_state_770_x) && !std::isnan(ideal_state_9410_x) && !std::isnan(ideal_state_10000_x) && std::abs(match_chi2) < 6"
    print("Run selection cuts =", selection)
    selected_b = preselection(
        cuts=selection,
        input_file="data/tracking_losses_ntuple_B_def.root",
    )
    selected_bj = preselection(
        cuts=selection,
        input_file="data/tracking_losses_ntuple_BJpsi_def.root",
    )
    merge_cmd = ["hadd", "-fk", selected, selected_b, selected_bj]
    print("Concatenate polarities ...")
    subprocess.run(merge_cmd, check=True)


cpp_files = []
if args.field_params:
    print("Parameterise magnet kink position ...")
    cpp_files.append(parameterise_magnet_kink(input_file=selected))
    print("Parameterise track model ...")
    cpp_files.append(parameterise_track_model(input_file=selected))

# selected_all_p = (
#     "nn_neural_net_training/data/param_data_B_default_thesis_selected_all_p.root"
# )
# if args.prepare_params_data:
#     selection_all_momenta = "chi2_comb < 5 && pid != 11"
#     print()
#     print("Run selection cuts =", selection_all_momenta)
#     selected_md_all_p = preselection(
#         cuts=selection_all_momenta,
#         outfile_postfix="selected_all_p",
#         input_file="nn_neural_net_training/data/param_data_B_default_thesis.root",
#     )

# if args.field_params:
#     print("Parameterise magnet kink position ...")
#     cpp_files.append(parameterise_magnet_kink(input_file=selected_all_p))
#     print("Parameterise track model ...")
#     cpp_files.append(parameterise_track_model(input_file=selected_all_p))
#     print("Parameterise search window ...")
#     cpp_files.append(parameterise_search_window(input_file=selected_all_p))
#     print("Parameterise magnetic field integral ...")
#     cpp_files.append(parameterise_field_integral(input_file=selected_all_p))
#     print("Parameterise Hough histogram binning ...")
#     cpp_files.append(parameterise_hough_histogram(input_file=selected_all_p))

###>>>
ghost_data = "neural_net_training/data/ghost_data.root"
if args.prepare_weights_data:
    merge_cmd = [
        "hadd",
        "-fk",
        ghost_data,
        "data/ghost_data_MD.root",
        "data/ghost_data_MU.root",
    ]
    print("Concatenate polarities for neural network training ...")
    subprocess.run(merge_cmd, check=True)
###<<<
if args.forward_weights:
    train_default_forward_ghost_mlp(prepare_data=args.prepare_weights_data)
    # FIXME: use env variable instead
    os.chdir(os.path.dirname(os.path.realpath(__file__)))
    train_veloUT_forward_ghost_mlp(prepare_data=args.prepare_weights_data)
    # this ensures that the directory is correct
    os.chdir(os.path.dirname(os.path.realpath(__file__)))
    cpp_files += parse_tmva_matrix_to_array(
        [
            "neural_net_training/result/GhostNNDataSet/weights/TMVAClassification_default_forward_ghost_mlp.class.C",
            "neural_net_training/result/GhostNNDataSet/weights/TMVAClassification_veloUT_forward_ghost_mlp.class.C",
        ],
    )
###>>>
if args.matching_weights:
    os.chdir(os.path.dirname(os.path.realpath(__file__)))
    train_matching_ghost_mlp(
        prepare_data=args.prepare,
        input_file="data/ghost_data_B_default_phi_eta.root",
        tree_name="PrMatchNN_3e224c41.PrMCDebugMatchToolNN/MVAInputAndOutput",
        outdir="neural_net_training",
        exclude_electrons=False,
        only_electrons=True,
    )
    # this ensures that the directory is correct
    os.chdir(os.path.dirname(os.path.realpath(__file__)))
    cpp_files += parse_tmva_matrix_to_array(
        [
            "neural_net_training/result/MatchNNDataSet/weights/TMVAClassification_matching_mlp.class.C",
        ],
        simd_type=True,
    )
###<<<
for file in cpp_files:
    subprocess.run(
        [
            "clang-format",
            "-i",
            f"{file}",
        ],
    )