{
 "cells": [
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-14T08:09:47.447782Z",
     "start_time": "2025-04-14T08:09:37.767746Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "import sys\n",
    "import os.path as osp\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "from src.dataset.dataset import SimpleIterDataset, EventDataset\n",
    "from src.utils.utils import to_filelist\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "matplotlib.rc('font', size=13)\n",
    "from src.plotting.plot_event import plot_event_comparison\n",
    "from src.dataset.functions_data import concat_events\n",
    "from src.utils.paths import get_path\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv()"
   ],
   "id": "6bae9707acf4a848",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-14T08:09:48.837368Z",
     "start_time": "2025-04-14T08:09:48.832767Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def remove_from_list(lst):\n",
    "    out = []\n",
    "    for item in lst:\n",
    "        if item in [\"hgcal\", \"data.txt\", \"test_file.root\"]:\n",
    "            continue\n",
    "        out.append(item)\n",
    "    return out\n",
    "\n",
    "#path = \"/eos/user/g/gkrzmanc/jetclustering/data/SVJ_std_UL2018_scouting_test_large/SVJ_mMed-700GeV_mDark-20GeV_rinv-0.7_alpha-peak\"\n",
    "def get_iter(path_to_ds):\n",
    "    return iter(EventDataset.from_directory(path_to_ds))\n"
   ],
   "id": "e7a7ef680143801e",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-14T08:09:49.693906Z",
     "start_time": "2025-04-14T08:09:49.608366Z"
    }
   },
   "cell_type": "code",
   "source": "dataset = get_iter(get_path(\"Feb26_2025_E1000_N500_noPartonFilter_GluonFixF/PFNano_s-channel_mMed-700_mDark-20_rinv-0.3_alpha-peak_13TeV-pythia8_n-1000\", \"preprocessed_data\"))",
   "id": "1549361c5b028634",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "get_pfcands_key\n"
     ]
    }
   ],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-14T08:09:51.053269Z",
     "start_time": "2025-04-14T08:09:50.666358Z"
    }
   },
   "cell_type": "code",
   "source": "e = next(dataset)",
   "id": "e0d491f2943f20e9",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-14T08:09:51.767492Z",
     "start_time": "2025-04-14T08:09:51.755352Z"
    }
   },
   "cell_type": "code",
   "source": [
    "random_generator = np.random.RandomState(seed=3)\n",
    "spl = EventDataset.pfcands_split_particles(e.pfcands, random_generator)"
   ],
   "id": "87c6ab0ccf50fa58",
   "outputs": [],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-14T08:09:52.738157Z",
     "start_time": "2025-04-14T08:09:52.725056Z"
    }
   },
   "cell_type": "code",
   "source": "spl.pt",
   "id": "6919a7adaa7a6376",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([44.3125,  5.3385, 15.2578, 13.7188,  9.7969,  7.9609,  7.3008,  7.0039,\n",
       "         6.9883,  6.7812,  5.3281,  4.8477,  4.6211,  4.5586,  4.4844,  4.4727,\n",
       "         4.2148,  3.7598,  3.2539,  3.2305,  3.1270,  3.0469,  3.0195,  3.0156,\n",
       "         2.7480,  2.6699,  2.5801,  2.3555,  2.3496,  2.2910,  2.1953,  2.1738,\n",
       "         2.0859,  2.0742,  2.0020,  1.9785,  1.9619,  1.9199,  1.8711,  1.8574,\n",
       "         1.8320,  1.7842,  1.7832,  1.7764,  1.7617,  1.6924,  1.6641,  1.6318,\n",
       "         1.6230,  1.6152,  1.6143,  1.5576,  1.5547,  1.5449,  1.5283,  1.4854,\n",
       "         1.4404,  1.4287,  1.3828,  1.3701,  1.3682,  1.3252,  1.2891,  1.2490,\n",
       "         1.2451,  1.2383,  1.2256,  1.2217,  1.2217,  1.1807,  1.1719,  1.1416,\n",
       "         1.1094,  1.1025,  1.1016,  1.0830,  1.0801,  1.0742,  1.0732,  1.0713,\n",
       "         1.0625,  1.0557,  1.0410,  1.0254,  1.0186,  1.0137,  1.0107,  0.9756,\n",
       "         0.9731,  0.9683,  0.9668,  0.9463,  0.9370,  0.9360,  0.9287,  0.9072,\n",
       "         0.9053,  0.8984,  0.8760,  0.8755,  0.8706,  0.8677,  0.8647,  0.8511,\n",
       "         0.8472,  0.8438,  0.8354,  0.8286,  0.8267,  0.8247,  0.8193,  0.8062,\n",
       "         0.8057,  0.7949,  0.7822,  0.7739,  0.7715,  0.7520,  0.7490,  0.7461,\n",
       "         0.7368,  0.7354,  0.7310,  0.7275,  0.7246,  0.7178,  0.7129,  0.7109,\n",
       "         0.7070,  0.7070,  0.7031,  0.6997,  0.6865,  0.6802,  0.6802,  0.6655,\n",
       "         0.6572,  0.6562,  0.6514,  0.6504,  0.6475,  0.6411,  0.6396,  0.6367,\n",
       "         0.6353,  0.6338,  0.6211,  0.6162,  0.6147,  0.6138,  0.6108,  0.6084,\n",
       "         0.6055,  0.6050,  0.6016,  0.6016, 44.3125,  5.3385,  5.3385],\n",
       "       dtype=torch.float64)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-14T08:09:54.479093Z",
     "start_time": "2025-04-14T08:09:54.468976Z"
    }
   },
   "cell_type": "code",
   "source": "e.pfcands.pt",
   "id": "2d7e0400c013c243",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([88.6250, 16.0156, 15.2578, 13.7188,  9.7969,  7.9609,  7.3008,  7.0039,\n",
       "         6.9883,  6.7812,  5.3281,  4.8477,  4.6211,  4.5586,  4.4844,  4.4727,\n",
       "         4.2148,  3.7598,  3.2539,  3.2305,  3.1270,  3.0469,  3.0195,  3.0156,\n",
       "         2.7480,  2.6699,  2.5801,  2.3555,  2.3496,  2.2910,  2.1953,  2.1738,\n",
       "         2.0859,  2.0742,  2.0020,  1.9785,  1.9619,  1.9199,  1.8711,  1.8574,\n",
       "         1.8320,  1.7842,  1.7832,  1.7764,  1.7617,  1.6924,  1.6641,  1.6318,\n",
       "         1.6230,  1.6152,  1.6143,  1.5576,  1.5547,  1.5449,  1.5283,  1.4854,\n",
       "         1.4404,  1.4287,  1.3828,  1.3701,  1.3682,  1.3252,  1.2891,  1.2490,\n",
       "         1.2451,  1.2383,  1.2256,  1.2217,  1.2217,  1.1807,  1.1719,  1.1416,\n",
       "         1.1094,  1.1025,  1.1016,  1.0830,  1.0801,  1.0742,  1.0732,  1.0713,\n",
       "         1.0625,  1.0557,  1.0410,  1.0254,  1.0186,  1.0137,  1.0107,  0.9756,\n",
       "         0.9731,  0.9683,  0.9668,  0.9463,  0.9370,  0.9360,  0.9287,  0.9072,\n",
       "         0.9053,  0.8984,  0.8760,  0.8755,  0.8706,  0.8677,  0.8647,  0.8511,\n",
       "         0.8472,  0.8438,  0.8354,  0.8286,  0.8267,  0.8247,  0.8193,  0.8062,\n",
       "         0.8057,  0.7949,  0.7822,  0.7739,  0.7715,  0.7520,  0.7490,  0.7461,\n",
       "         0.7368,  0.7354,  0.7310,  0.7275,  0.7246,  0.7178,  0.7129,  0.7109,\n",
       "         0.7070,  0.7070,  0.7031,  0.6997,  0.6865,  0.6802,  0.6802,  0.6655,\n",
       "         0.6572,  0.6562,  0.6514,  0.6504,  0.6475,  0.6411,  0.6396,  0.6367,\n",
       "         0.6353,  0.6338,  0.6211,  0.6162,  0.6147,  0.6138,  0.6108,  0.6084,\n",
       "         0.6055,  0.6050,  0.6016,  0.6016], dtype=torch.float64)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 7
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "",
   "id": "8f771993addd4bc1"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}