Scripts/Optimization/Optuna_tuning.py

0001 #!/usr/bin/env python3
0002 import sys
0003 import optuna
0004 import logging
0005 import uproot
0006 import matplotlib
0007 import subprocess
0008 import json
0009 import pandas as pd
0010 from pathlib import Path
0011
0012 matplotlib.use("pdf")
0013
0014 srcDir = Path(__file__).resolve().parent
0015
0016
0017 def run_ckf(params, names, outDir):
0018     if len(params) != len(names):
0019         raise Exception("Length of Params must equal names")
0020
0021     ckf_script = srcDir / "ckf.py"
0022     nevts = "--nEvents=1"
0023     indir = "--indir=" + str(srcDir)
0024     outdir = "--output=" + str(outDir)
0025
0026     ret = ["python"]
0027     ret.append(ckf_script)
0028     ret.append(nevts)
0029     ret.append(indir)
0030     ret.append(outdir)
0031
0032     i = 0
0033     for param in params:
0034         arg = "--sf_" + names[i] + "=" + str(param)
0035         ret.append(arg)
0036         i += 1
0037
0038     # Run CKF for the given parameters
0039     subprocess.call(ret)
0040
0041
0042 class Objective:
0043     def __init__(self, k_dup, k_time):
0044         self.res = {
0045             "eff": [],
0046             "fakerate": [],
0047             "duplicaterate": [],
0048             "runtime": [],
0049         }
0050
0051         self.k_dup = k_dup
0052         self.k_time = k_time
0053
0054     def __call__(self, trial, ckf_perf=True):
0055         params = []
0056
0057         maxSeedsPerSpM = trial.suggest_int("maxSeedsPerSpM", 0, 10)
0058         params.append(maxSeedsPerSpM)
0059         cotThetaMax = trial.suggest_float("cotThetaMax", 5.0, 10.0)
0060         params.append(cotThetaMax)
0061         sigmaScattering = trial.suggest_float("sigmaScattering", 0.2, 50)
0062         params.append(sigmaScattering)
0063         radLengthPerSeed = trial.suggest_float("radLengthPerSeed", 0.001, 0.1)
0064         params.append(radLengthPerSeed)
0065         impactMax = trial.suggest_float("impactMax", 0.1, 25)
0066         params.append(impactMax)
0067         maxPtScattering = trial.suggest_float("maxPtScattering", 1, 50)
0068         params.append(maxPtScattering)
0069         deltaRMin = trial.suggest_float("deltaRMin", 0.25, 30)
0070         params.append(deltaRMin)
0071         deltaRMax = trial.suggest_float("deltaRMax", 50, 300)
0072         params.append(deltaRMax)
0073         keys = [
0074             "maxSeedsPerSpM",
0075             "cotThetaMax",
0076             "sigmaScattering",
0077             "radLengthPerSeed",
0078             "impactMax",
0079             "maxPtScattering",
0080             "deltaRMin",
0081             "deltaRMax",
0082         ]
0083
0084         get_tracking_perf(self, ckf_perf, params, keys)
0085
0086         efficiency = self.res["eff"][-1]
0087         penalty = (
0088             self.res["fakerate"][-1]
0089             + self.res["duplicaterate"][-1] / self.k_dup
0090             + self.res["runtime"][-1] / self.k_time
0091         )
0092
0093         return efficiency - penalty
0094
0095
0096 def get_tracking_perf(self, ckf_perf, params, keys):
0097     if ckf_perf:
0098         outDirName = "Output_CKF"
0099         outputfile = srcDir / outDirName / "performance_ckf.root"
0100         effContName = "particles"
0101         contName = "tracks"
0102     else:
0103         outDirName = "Output_Seeding"
0104         outputfile = srcDir / outDirName / "performance_seeding.root"
0105         effContName = "seeds"
0106         contName = "seeds"
0107
0108     outputDir = Path(srcDir / outDirName)
0109     outputDir.mkdir(exist_ok=True)
0110     run_ckf(params, keys, outputDir)
0111     rootFile = uproot.open(outputfile)
0112     self.res["eff"].append(rootFile["eff_" + effContName].member("fElements")[0])
0113     self.res["fakerate"].append(rootFile["fakerate_" + contName].member("fElements")[0])
0114     self.res["duplicaterate"].append(
0115         rootFile["duplicaterate_" + contName].member("fElements")[0]
0116     )
0117
0118     timingfile = srcDir / outDirName / "timing.tsv"
0119     timing = pd.read_csv(timingfile, sep="\t")
0120
0121     if ckf_perf:
0122         time_ckf = float(
0123             timing[timing["identifier"].str.match("Algorithm:TrackFindingAlgorithm")][
0124                 "time_perevent_s"
0125             ]
0126         )
0127
0128     time_seeding = float(
0129         timing[timing["identifier"].str.match("Algorithm:SeedingAlgorithm")][
0130             "time_perevent_s"
0131         ]
0132     )
0133     if ckf_perf:
0134         self.res["runtime"].append(time_ckf + time_seeding)
0135     else:
0136         self.res["runtime"].append(time_seeding)
0137
0138
0139 def main():
0140     k_dup = 5
0141     k_time = 5
0142
0143     # Initializing the objective (score) function
0144     objective = Objective(k_dup, k_time)
0145
0146     """
0147     start_values = {
0148         "maxSeedsPerSpM": 1,
0149         "cotThetaMax": 7.40627,
0150         "sigmaScattering": 50,
0151         "radLengthPerSeed": 0.1,
0152         "impactMax": 3.0,
0153         "maxPtScattering": 10.0,
0154         "deltaRMin": 1.0,
0155         "deltaRMax": 60.0
0156 }
0157
0158     """
0159     # Optuna logger
0160     optuna.logging.get_logger("optuna").addHandler(logging.StreamHandler(sys.stdout))
0161     study_name = "test_study"
0162     storage_name = "sqlite:///{}.db".format(study_name)
0163
0164     # creating a new optuna study
0165     study = optuna.create_study(
0166         study_name=study_name,
0167         storage=storage_name,
0168         direction="maximize",
0169         load_if_exists=True,
0170     )
0171
0172     # study.enqueue_trial(start_values)
0173     # Start Optimization
0174     study.optimize(objective, n_trials=3)
0175
0176     # Printout the best trial values
0177     print("Best Trial until now", flush=True)
0178     for key, value in study.best_trial.params.items():
0179         print(f"    {key}: {value}", flush=True)
0180
0181     outputDir = Path("OptunaResults")
0182     outputDir.mkdir(exist_ok=True)
0183
0184     with open(outputDir / "results.json", "w") as fp:
0185         json.dump(study.best_params, fp)
0186
0187
0188 if __name__ == "__main__":
0189     main()