refactor

LuisFelipeCoelho · Jan 19, 2024 · 0819498 · 0819498
1 parent 8d986ff
commit 0819498
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 84 deletions.
diff --git a/Examples/Scripts/Optimization/Optuna_tuning.py b/Examples/Scripts/Optimization/Optuna_tuning.py
@@ -109,48 +109,7 @@ def __call__(self, trial, ckf_perf=True):
             "deltaRMax",
         ]
 
-        if ckf_perf:
-            outDirName = "Output_CKF"
-            outputfile = srcDir / outDirName / "performance_ckf.root"
-            effContName = "particles"
-            contName = "tracks"
-        else:
-            outDirName = "Output_Seeding"
-            outputfile = srcDir / outDirName / "performance_seeding.root"
-            effContName = "seeds"
-            contName = "seeds"
-
-        outputDir = Path(srcDir / outDirName)
-        outputDir.mkdir(exist_ok=True)
-        run_ckf(params, keys, outputDir)
-        rootFile = uproot.open(outputfile)
-        self.res["eff"].append(rootFile["eff_" + effContName].member("fElements")[0])
-        self.res["fakerate"].append(
-            rootFile["fakerate_" + contName].member("fElements")[0]
-        )
-        self.res["duplicaterate"].append(
-            rootFile["duplicaterate_" + contName].member("fElements")[0]
-        )
-
-        timingfile = srcDir / outDirName / "timing.tsv"
-        timing = pd.read_csv(timingfile, sep="\t")
-
-        if ckf_perf:
-            time_ckf = float(
-                timing[
-                    timing["identifier"].str.match("Algorithm:TrackFindingAlgorithm")
-                ]["time_perevent_s"]
-            )
-
-        time_seeding = float(
-            timing[timing["identifier"].str.match("Algorithm:SeedingAlgorithm")][
-                "time_perevent_s"
-            ]
-        )
-        if ckf_perf:
-            self.res["runtime"].append(time_ckf + time_seeding)
-        else:
-            self.res["runtime"].append(time_seeding)
+        get_tracking_perf(self, ckf_perf, params, keys)
 
         efficiency = self.res["eff"][-1]
         penalty = (
@@ -162,6 +121,49 @@ def __call__(self, trial, ckf_perf=True):
         return efficiency - penalty
 
 
+def get_tracking_perf(self, ckf_perf, params, keys):
+    if ckf_perf:
+        outDirName = "Output_CKF"
+        outputfile = srcDir / outDirName / "performance_ckf.root"
+        effContName = "particles"
+        contName = "tracks"
+    else:
+        outDirName = "Output_Seeding"
+        outputfile = srcDir / outDirName / "performance_seeding.root"
+        effContName = "seeds"
+        contName = "seeds"
+
+    outputDir = Path(srcDir / outDirName)
+    outputDir.mkdir(exist_ok=True)
+    run_ckf(params, keys, outputDir)
+    rootFile = uproot.open(outputfile)
+    self.res["eff"].append(rootFile["eff_" + effContName].member("fElements")[0])
+    self.res["fakerate"].append(rootFile["fakerate_" + contName].member("fElements")[0])
+    self.res["duplicaterate"].append(
+        rootFile["duplicaterate_" + contName].member("fElements")[0]
+    )
+
+    timingfile = srcDir / outDirName / "timing.tsv"
+    timing = pd.read_csv(timingfile, sep="\t")
+
+    if ckf_perf:
+        time_ckf = float(
+            timing[timing["identifier"].str.match("Algorithm:TrackFindingAlgorithm")][
+                "time_perevent_s"
+            ]
+        )
+
+    time_seeding = float(
+        timing[timing["identifier"].str.match("Algorithm:SeedingAlgorithm")][
+            "time_perevent_s"
+        ]
+    )
+    if ckf_perf:
+        self.res["runtime"].append(time_ckf + time_seeding)
+    else:
+        self.res["runtime"].append(time_seeding)
+
+
 def main():
     k_dup = 5
     k_time = 5

diff --git a/Examples/Scripts/Optimization/Orion_tuning.py b/Examples/Scripts/Optimization/Orion_tuning.py
@@ -32,6 +32,8 @@
 
 from orion.client import build_experiment
 
+from Optuna_tuning import get_tracking_perf, run_ckf
+
 srcDir = Path(__file__).resolve().parent
 
 
@@ -105,48 +107,7 @@ def __call__(
             "deltaRMax",
         ]
 
-        if ckf_perf:
-            outDirName = "Output_CKF"
-            outputfile = srcDir / outDirName / "performance_ckf.root"
-            effContName = "particles"
-            contName = "tracks"
-        else:
-            outDirName = "Output_Seeding"
-            outputfile = srcDir / outDirName / "performance_seeding.root"
-            effContName = "seeds"
-            contName = "seeds"
-
-        outputDir = Path(srcDir / outDirName)
-        outputDir.mkdir(exist_ok=True)
-        run_ckf(params, keys, outputDir)
-        rootFile = uproot.open(outputfile)
-        self.res["eff"].append(rootFile["eff_" + effContName].member("fElements")[0])
-        self.res["fakerate"].append(
-            rootFile["fakerate_" + contName].member("fElements")[0]
-        )
-        self.res["duplicaterate"].append(
-            rootFile["duplicaterate_" + contName].member("fElements")[0]
-        )
-
-        timingfile = srcDir / outDirName / "timing.tsv"
-        timing = pd.read_csv(timingfile, sep="\t")
-
-        if ckf_perf:
-            time_ckf = float(
-                timing[
-                    timing["identifier"].str.match("Algorithm:TrackFindingAlgorithm")
-                ]["time_perevent_s"]
-            )
-
-        time_seeding = float(
-            timing[timing["identifier"].str.match("Algorithm:SeedingAlgorithm")][
-                "time_perevent_s"
-            ]
-        )
-        if ckf_perf:
-            self.res["runtime"].append(time_ckf + time_seeding)
-        else:
-            self.res["runtime"].append(time_seeding)
+        get_tracking_perf(self, ckf_perf, params, keys)
 
         efficiency = self.res["eff"][-1]
         penalty = (