tracking-parametrisation-tuner/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.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",
    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 = "neural_net_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 = "neural_net_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 = "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=False,
        input_file="data/ghost_data_B_electron_weights.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}",
        ],
    )