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 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_selected.root"
if args.prepare_params_data:
    selection = "chi2_comb < 5 && pt > 10 && p > 1500 && p < 100000 && pid != 11"
    print("Run selection cuts =", selection)
    selected_md = preselection(
        cuts=selection,
        input_file="data/param_data_MD.root",
    )
    selected_mu = preselection(
        cuts=selection,
        input_file="data/param_data_MU.root",
    )
    merge_cmd = ["hadd", "-fk", selected, selected_md, selected_mu]
    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_electron_training/data/param_data_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="data/param_data_MD.root",
    )
    selected_mu_all_p = preselection(
        cuts=selection_all_momenta,
        outfile_postfix="selected_all_p",
        input_file="data/param_data_MU.root",
    )
    merge_cmd = ["hadd", "-fk", selected_all_p, selected_md_all_p, selected_mu_all_p]
    print("Concatenate polarities ...")
    subprocess.run(merge_cmd, check=True)

if args.field_params:
    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 = "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 and not args.residuals:
    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=100e3,
        n_train_bkg=100e3,
        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}",
        ],
    )