tracking-parametrisation-tuner/electron_main.py

# flake8: noqaq
import os
import subprocess
import argparse
from parameterisations.parameterise_magnet_kink 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 import preselection
from parameterisations.train_forward_ghost_mlps import (
    train_default_forward_ghost_mlp,
    train_veloUT_forward_ghost_mlp,
)
from parameterisations.residual_train_matching_ghost_mlps_electron import (
    res_train_matching_ghost_mlp,
)
from parameterisations.train_matching_ghost_mlps_electron import (
    train_matching_ghost_mlp,
)

from parameterisations.utils.parse_tmva_matrix_to_array_electron 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(
#     "-r",
#     "--residuals",
#     action="store_true",
#     help="Trains neural network with residual tracks.",
# )
parser.add_argument(
    "-p",
    "--prepare",
    action="store_true",
    default=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 = "nn_electron_training/data/param_data_B_default_thesis_selected.root"
if args.prepare and args.field_params:
    selection: str = "isElectron == 1"
    print("Selection Cuts = ", selection)
    selected_sample = preselection(
        cuts=selection,
        input_file="nn_electron_training/data/param_data_B_default_thesis.root",
    )

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))

ghost_data = "data/ghost_data.root"
if args.prepare_weights_data:
    merge_cmd = [
        "hadd",
        "-fk",
        ghost_data,
        "data/ghost_data_B.root",
        "data/ghost_data_D.root",
    ]
    print("Concatenate decays 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(
        [
            "nn_electron_training/result/GhostNNDataSet/weights/TMVAClassification_default_forward_ghost_mlp.class.C",
            "nn_electron_training/result/GhostNNDataSet/weights/TMVAClassification_veloUT_forward_ghost_mlp.class.C",
        ],
    )

# if args.matching_weights and args.residuals:
#     os.chdir(os.path.dirname(os.path.realpath(__file__)))
#     res_train_matching_ghost_mlp(
#         prepare_data=args.prepare,
#         input_file="data/ghost_data_B_default_only_e_as_seed.root",
#         tree_name="PrMatchNN_b60a058d.PrMCDebugMatchToolNN/MVAInputAndOutput",  # e6feac0d, B: 3e224c41, B res: 1e13cc7e, D: 8cb154ca
#         exclude_electrons=False,
#         only_electrons=True,
#         residuals="PrMatchNN_1e13cc7e.PrMCDebugMatchToolNN/MVAInputAndOutput",
#         outdir="nn_electron_training",
#         n_train_signal=0,
#         n_train_bkg=20e3,
#         n_test_signal=1e3,
#         n_test_bkg=5e3,
#     )
#     # this ensures that the directory is correct
#     os.chdir(os.path.dirname(os.path.realpath(__file__)))
#     cpp_files += parse_tmva_matrix_to_array(
#         [
#             "nn_electron_training/result/MatchNNDataSet/weights/TMVAClassification_matching_mlp.class.C",
#         ],
#         simd_type=True,
#     )

file_name = "seed"

tree_names = {}
tree_names["seed"] = "PrMatchNN_b60a058d.PrMCDebugMatchToolNN/MVAInputAndOutput"
tree_names["def"] = "PrMatchNN.PrMCDebugMatchToolNN/MVAInputAndOutput"

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_vars_thesis.root",
        tree_name=tree_names[file_name],
        exclude_electrons=False,
        only_electrons=True,
        filter_seeds=True,
        outdir="nn_electron_training",
        n_train_signal=150e3,
        n_train_bkg=150e3,
        n_test_signal=10e3,
        n_test_bkg=10e3,
    )
    # this ensures that the directory is correct
    os.chdir(os.path.dirname(os.path.realpath(__file__)))
    cpp_files += parse_tmva_matrix_to_array(
        [
            "nn_electron_training/result/MatchNNDataSet/weights/TMVAClassification_matching_mlp.class.C",
        ],
        simd_type=True,
    )

for file in cpp_files:
    subprocess.run(
        [
            "clang-format",
            "-i",
            f"{file}",
        ],
    )