ml_pid_cbm.validate_multiple_models

 1import argparse
 2import os
 3import sys
 4from concurrent.futures import ThreadPoolExecutor
 5from shutil import copy2
 6from typing import List, Set
 7
 8import pandas as pd
 9
10from validate_model import ValidateModel
11
12
13class ValidateMultipleModels(ValidateModel):
14    """
15    Class for validating data from multiple models.
16    Inherits from ValidateModel
17    """
18
19    def __init__(self, json_file_name: str, files_list: Set[str], n_workers: int = 1):
20        super().__init__(-12, 12, False, json_file_name, None)
21        self.particles_df = self.load_pickles(files_list, n_workers)
22
23    @staticmethod
24    def load_pickles(files_list: Set[str], n_workers: int = 1) -> pd.DataFrame:
25        """Loads multiple pickle files produced by validate_model module.
26
27        Args:
28            files_list (Set[str]): Files list containg picle files with datasets.
29            n_workers (int, optional): Number of workers for multithreading. Defaults to 1.
30
31        Returns:
32            pd.DataFrame: Dataframe with merged datasets.
33        """
34        with ThreadPoolExecutor(max_workers=n_workers) as executor:
35            results = list(executor.map(pd.read_pickle, files_list))
36            whole_df = pd.concat(results, ignore_index=True)
37        return whole_df
38
39
40def parse_args(args: List[str]) -> argparse.Namespace:
41    """
42    Arguments parser for the main method.
43
44    Args:
45        args (List[str]): Arguments from the command line, should be sys.argv[1:].
46
47    Returns:
48        argparse.Namespace: argparse.Namespace containg args
49    """
50    parser = argparse.ArgumentParser(
51        prog="ML_PID_CBM ValidateMultipleModels",
52        description="Program for loading multiple validated PID ML models",
53    )
54    parser.add_argument(
55        "--modelnames",
56        "-m",
57        nargs="+",
58        required=True,
59        type=str,
60        help="Names of folders containing trained and validated ML models.",
61    )
62    parser.add_argument(
63        "--config",
64        "-c",
65        nargs=1,
66        required=True,
67        type=str,
68        help="Filename of path of config json file.",
69    )
70    parser.add_argument(
71        "--nworkers",
72        "-n",
73        type=int,
74        default=1,
75        help="Max number of workers for ThreadPoolExecutor which reads Root tree with data.",
76    )
77    return parser.parse_args(args)
78
79
80if __name__ == "__main__":
81    args = parse_args(sys.argv[1:])
82    # config  arguments to be loaded from args
83    json_file_name = args.config[0]
84    models = args.modelnames
85    n_workers = args.nworkers
86    pickle_files = {f"{model}/validated_data.pickle" for model in models}
87    validate = ValidateMultipleModels(json_file_name, pickle_files, n_workers)
88    # new folder for all files
89    json_file_path = os.path.join(os.getcwd(), json_file_name)
90    if not os.path.exists("all_models"):
91        os.makedirs("all_models")
92    os.chdir("all_models")
93    copy2(json_file_path, os.getcwd())
94    # graphs
95    validate.confusion_matrix_and_stats()
96    print("Generating plots...")
97    validate.generate_plots()