Add seeds and deterministic thread-getting to make this reproducible

Author: alexandertuna

RecoTracker/LSTCore/standalone/analysis/DNN/embed_train.ipynb

@@ -19,6 +19,15 @@
     "import awkward as ak # Using awkward array for easier handling of jagged data\n",
     "import time # For timing steps\n",
     "\n",
+    "# Set seeds for reproducibility\n",
+    "seed_value = 42\n",
+    "random.seed(seed_value)\n",
+    "np.random.seed(seed_value)\n",
+    "torch.manual_seed(seed_value)\n",
+    "if torch.cuda.is_available():\n",
+    "    torch.cuda.manual_seed(seed_value)\n",
+    "os.environ['PYTHONHASHSEED'] = str(seed_value)\n",
+    "\n",
     "def load_root_file(file_path, branches=None, print_branches=False):\n",
     "    all_branches = {}\n",
     "    with uproot.open(file_path) as file:\n",
@@ -381,6 +390,7 @@
     "    dis_pairs = idxs_triu[dis_mask]\n",
     "\n",
     "    # down-sample\n",
+    "    random.seed(evt_idx)\n",
     "    if len(sim_pairs) > max_sim:\n",
     "        sim_pairs = sim_pairs[random.sample(range(len(sim_pairs)), max_sim)]\n",
     "    if len(dis_pairs) > max_dis:\n",
@@ -417,7 +427,7 @@
     "\n",
     "    with ProcessPoolExecutor(max_workers=n_workers) as pool:\n",
     "        futures = [pool.submit(_pairs_single_event, *args) for args in work_args]\n",
-    "        for fut in as_completed(futures):\n",
+    "        for fut in futures:\n",
     "            evt_idx, sim_pairs_evt, dis_pairs_evt = fut.result()\n",
     "            F = features_per_event[evt_idx]\n",
     "            D = displaced_per_event[evt_idx]\n",
@@ -546,6 +556,7 @@
     "    dis_pairs = np.column_stack((idx_p[dis_mask], idx_t[dis_mask]))\n",
     "\n",
     "    # down-sample\n",
+    "    random.seed(evt_idx)\n",
     "    if len(sim_pairs) > max_sim:\n",
     "        sim_pairs = sim_pairs[random.sample(range(len(sim_pairs)), max_sim)]\n",
     "    if len(dis_pairs) > max_dis:\n",
@@ -584,7 +595,7 @@
     "        )\n",
     "        for ev in range(len(features_per_event))\n",
     "    ]\n",
-    "    for fut in as_completed(futures):\n",
+    "    for fut in futures:\n",
     "        _, packed = fut.result()\n",
     "        # accumulate\n",
     "        sim_evt = sum(1 for _,_,lbl,_ in packed if lbl == 0)\n",