move presets to a separate file & cleanup

PurdueAF · Mar 28, 2024 · 51cd724 · 51cd724
1 parent 828c149
commit 51cd724
Show file tree

Hide file tree

Showing 7 changed files with 51 additions and 71 deletions.
diff --git a/notebooks/2.1_Single-file.ipynb b/notebooks/2.1_Single-file.ipynb
diff --git a/scenarios/generator_2p1.py b/scenarios/generator_2p1.py
@@ -3,30 +3,7 @@
 import yaml
 import glob
 
-column_presets = {
-    # "full_event": {
-    #     # the bechmark will limit this to actual total number of columns
-    #     "method": "n_columns",
-    #     "values": 100000
-    # },
-    "main_collections": {
-        "method": "collections",
-        "values": ["Jet", "Photon", "Tau", "Electron", "Muon"]
-    },
-    "muons_only": {
-        "method": "collections",
-        "values": ["Muon"]
-    },
-    "hmm_columns": {
-        "method": "column_list",
-        "values": [
-            "run", "luminosityBlock", "HLT_IsoMu24", "PV_npvsGood", "fixedGridRhoFastjetAll",
-            "Muon_pt", "Muon_eta", "Muon_phi", "Muon_mass", "Muon_charge", "Muon_pfRelIso04_all", "Muon_mediumId", "Muon_ptErr",
-            "Electron_pt", "Electron_eta", "Electron_mvaFall17V2Iso_WP90",
-            "Jet_pt", "Jet_eta", "Jet_phi", "Jet_mass",
-        ]
-    }
-}
+from scenarios.presets import column_presets
 
 
 default_config = {
@@ -69,6 +46,8 @@ def generate_configs(save_dir="./"):
             config = copy.deepcopy(default_config)
             config["executor"]["n_workers"] = n_workers
             config["processor"]["columns"] = column_setup
+
+            # Custom labels to save to output dataframe
             config["custom_labels"] = {
                 "column_setup": label
             }

diff --git a/src/benchmark.py b/src/benchmark.py
@@ -5,13 +5,13 @@
 import tqdm
 import pandas as pd
 
-from time_profiler import time_profiler as tp
-from data_loader import get_file_list
-from uproot_processor import UprootProcessor
+from src.time_profiler import time_profiler as tp
+from src.data_loader import get_file_list
+from src.uproot_processor import UprootProcessor
 
-from executors.sequential import SequentialExecutor
-from executors.futures import FuturesExecutor
-from executors.dask import DaskLocalExecutor, DaskGatewayExecutor
+from src.executors.sequential import SequentialExecutor
+from src.executors.futures import FuturesExecutor
+from src.executors.dask import DaskLocalExecutor, DaskGatewayExecutor
 executors = {
     'sequential': SequentialExecutor,
     'futures': FuturesExecutor,

diff --git a/src/executors/dask.py b/src/executors/dask.py
@@ -1,4 +1,4 @@
-from executors.base import BaseExecutor
+from src.executors.base import BaseExecutor
 import dask
 from dask.distributed import LocalCluster, Client
 from dask_gateway import Gateway

diff --git a/src/executors/futures.py b/src/executors/futures.py
@@ -1,4 +1,4 @@
-from executors.base import BaseExecutor
+from src.executors.base import BaseExecutor
 from concurrent import futures
 
 class FuturesExecutor(BaseExecutor):

diff --git a/src/executors/sequential.py b/src/executors/sequential.py
@@ -1,4 +1,4 @@
-from executors.base import BaseExecutor
+from src.executors.base import BaseExecutor
 
 
 class SequentialExecutor(BaseExecutor):

diff --git a/src/uproot_processor.py b/src/uproot_processor.py
@@ -1,4 +1,4 @@
-from time_profiler import time_profiler as tp
+from src.time_profiler import time_profiler as tp
 import pandas as pd
 import numpy as np
 import uproot