example_event_display.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import uproot4\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "from matplotlib.collections import PatchCollection\n",
    "from matplotlib import cm\n",
    "from matplotlib.colors import LogNorm, Normalize\n",
    "import awkward1 as ak\n",
    "from lxml import etree as ET\n",
    "import os\n",
    "from pathlib import Path\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import time\n",
    "\n",
    "from event_display import get_module_geometry, get_module_positions, plot_calorimeter_hits, build_calorimeter_section\n",
    "from geant3_parser import Geant3DataFile, build_train_set_xy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time open 0.35092759132385254\n",
      "Time split 0.2915961742401123\n",
      "Start preparing events...\n",
      "Total events prepare time = 0.16317129135131836\n",
      "max hit value = 0.7586345789791428\n",
      "max e = 5.4983\n",
      "Inputs shape original = (5000, 121)\n",
      "Inputs shape new = (5000, 11, 11, 1)\n"
     ]
    }
   ],
   "source": [
    "file_name = os.path.join('data', 'shower_geant3_new.dat')\n",
    "\n",
    "\n",
    "# file_name = 'sample_data.txt'\n",
    "data_file = Geant3DataFile(file_name, skip_lines=3)\n",
    "\n",
    "# split into input (X) and output (y) variables\n",
    "parse_start = time.time()\n",
    "print(f\"Start preparing events...\")\n",
    "\n",
    "\n",
    "cells, true_values = build_train_set_xy(data_file, 5000, normalize=True)\n",
    "parse_end = time.time()\n",
    "\n",
    "print(f\"Total events prepare time = {parse_end - parse_start}\")\n",
    "print(f\"max hit value = {np.max(cells)}\")\n",
    "print(f\"max e = {np.max(true_values[:,0])}\")\n",
    "print(f\"Inputs shape original = {np.shape(cells)}\")\n",
    "\n",
    "# Reshape it to 11x11x1\n",
    "cells = np.reshape(cells, (len(cells), 11, 11, 1))  # -1 => autodetermine\n",
    "\n",
    "# Pad with 1 row and column of zeroes, so it divides by 2\n",
    "# inputs = np.pad(inputs, ((0,0), (0,1), (0,1), (0,0)), mode='constant', constant_values=0)\n",
    "print(f\"Inputs shape new = {np.shape(cells)}\")\n",
    "\n",
    "\n",
    "energies = true_values[:, 0]        # True energy of the particle\n",
    "cell_sums = true_values[:, 1]       # Sum of all cells (you don't need it)\n",
    "incident_x = true_values[:, 2]      # X of incidence inside the cell \n",
    "incident_y = true_values[:, 3]      # Y "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def display_event(event_data, fig=None, ax=None, cal_size=11, cell_size=1):\n",
    "    # constants\n",
    "    size_x = cell_size\n",
    "    size_y = cell_size\n",
    "    dx = size_x / 2.0\n",
    "    dy = size_y / 2.0\n",
    "\n",
    "    # go through all cells and calculate their centers\n",
    "    centers = np.arange(-cal_size/2.0 + cell_size/2, cal_size/2 + cell_size/2, 1)\n",
    "    positions = []\n",
    "    for y_iter in range(cal_size):\n",
    "        for x_iter in range(cal_size):\n",
    "            positions.append((centers[x_iter], centers[y_iter]))\n",
    "\n",
    "    # plot calorimeter with empty cells\n",
    "    if not fig or not ax:\n",
    "        fig, ax = plt.subplots()\n",
    "    build_calorimeter_section(ax, positions, 1, 1)\n",
    "\n",
    "    # Create a heat map\n",
    "    norm = LogNorm()\n",
    "    norm.autoscale(event_data)\n",
    "    cmap = cm.get_cmap('inferno')\n",
    "\n",
    "    # Convert data to rectangular patches    \n",
    "    module_rects = []\n",
    "    for y_iter in range(cal_size):\n",
    "        for x_iter in range(cal_size):\n",
    "            x = centers[x_iter]\n",
    "            y = centers[y_iter]\n",
    "            weight = event_data[x_iter][y_iter][0]\n",
    "            #print(x,y,weight)\n",
    "            patch = patches.Rectangle((x-dx, y-dy), size_x, size_y, edgecolor='black', facecolor=cmap(norm(weight)))\n",
    "            module_rects.append(patch)\n",
    "\n",
    "    # plot rectangles with data\n",
    "    col = PatchCollection(module_rects, match_original=True)\n",
    "    ax.add_collection(col)\n",
    "\n",
    "    # plot heatmap legend\n",
    "    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)\n",
    "    sm.set_array(event_data)\n",
    "    fig.colorbar(sm)\n",
    "    \n",
    "    return fig, ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<Figure size 432x288 with 2 Axes>, <AxesSubplot:>)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAWsAAAD7CAYAAACsV7WPAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAATYklEQVR4nO3df6zddX3H8dfrthSHSlFa5kJbCw4ZTjGwC8OgcwN/dPLDJXOuLhqCzAY3EBONgsSMxSwuc/En6nYFnEYiIfiLbhXEIM4tgtwyQaGiFQcUZVzchr824Pa89sc5jaW77Tnt93N7Pt/veT7IN+m599v39317T97nw+f7+X7eTiIAQN2mxp0AAGA4ijUAtADFGgBagGINAC1AsQaAFqBYA0ALFCnWtg+xfY3t79jeYvsFJeICAPqWForzAUnXJXmV7WWSDioUFwAgyU0firF9sKTbJR2ZEYOtWLEia9eubXRdAJNh8+bNDydZ2STGy9cdmx8//LMRrvWD65Osa3KtxVJiZH2kpDlJH7f9fEmbJV2Q5Oc7n2R7g6QNkrRmzRrNzs4WuDSArrN9b9MYDz/8U938jb8cet6yJWetaHqtxVJiznqppOMlfTTJcZJ+LunCXU9KMpNkOsn0ypWNPiQBYC9FyfzQQ9Jy2zO2zxh3xrsqMbLeJmlbklsGr6/RAsUaAMYmUrJ9lDMfSbJhsdPZF42LdZIHbd9v++gkd0s6VdJdzVMDgDKiqNcfObdWqdUg50u6crAS5B5JZxeKCwAF9NTr/e+4k2ikSLFO8k1J0yViAUBxidJjZA0A9WMaBABqlx2rPVqLYg1gAkTqPT7uJBqZ+GI9NTUlWpsBT2RbvV5v3GkUkzCybr0kuuSSSxrHueSSS4hDnE7F6ZT0pHlWgwBA/RhZA0DdrMgs3QOAyiUSxRoAakexBoD6JfL2R0c5c7ntGUkbk2xc5Kz2CsUawAQYec66u7vuAUD9IvVG2iK1WhRrAN3HDUYAaAczsgaAyoVpEABogcjzj407iUYo1gA6zwnTIADQChRrAKgcI2sAaAmKNQDULnLLmyl4HF1SpqenMzs7u9+vuxDb404BqFItHZRsb04y3STGbx1zUL7+8WcPPe/AF9ze+FqLhZG1ynTFqLHTB3GI0yRO1zjtHllTrAF0Hw/FAEBLtHzOemrcCQDAokv6xXrYMdjP2vYZ4055V4ysAUyAyPOPj3Ji9/eztr1E0qykB5KcXiouADQWtX4apOTI+gJJWyQdXDAmABSQ1hfrInPWtldJOk3SZSXiAUBRUX81yLCjYqVG1u+X9DZJT93dCbY3SNogSWvWrCl0WQAYzh14grHxyNr26ZIeSrJ5T+clmUkynWR65cqVTS8LAHtntNUg1Soxsj5Z0pm2XyHpSZIOtv2pJK8tEBsAmkuk+Xb3YGw8sk5yUZJVSdZKWi/pRgo1gKpEUi/Dj4qxzhrAZKh8mmOYosU6yU2SbioZEwCaa//SPUbWALpvxzRIi1GsAUwGtkgFgMol0ny7izWdYugUAyyoS51ipo+Yyi1/sWzoeUvPfpROMTWrrUNHV+O8c835jeO8674Pqczbdr5YnFI/V22/ry5JWj8LQrEGMCG4wQgAlYskRtYA0AJ1b6o3FMUaQPdFSq/diwnowQhgMvRGOOjBCABjNtrIuvs9GAGgWh2YBqFYA5gAHnVkXS2KNYDui5Tt7b5FR7EGMBl6FGsAqFuYBgGAVkgo1gBQP6ZBAKBuYekeALSBWQ0CANWLWj8NQqcYOsXsF1Oyemr3fsIL6erPJXWrU8zxv3ZAvnb2iqHnPeXdD9Ippma1dego1+mjTCeU7/3ByY2jHPX5fy0W5+YXn9Y4zklf/afqfl+1xekWsxoEAKrXgWkQijWAzoukHjcYAaByMSNrAGiDtq+zbvxRY3u17a/Y3mL7TtsXlEgMAEpKPPSoWZnlAtJbktxm+6mSNtu+IcldBWIDQHNMg0hJfiTpR4M//9T2FkmHS6JYA6hG26dBis5Z214r6ThJtyzwvQ2SNkjSmjVrSl4WAPYoaf9qkGLZ236KpM9IenOSn+z6/SQzSaaTTK9cubLUZQFgBFYyNfSoWZGRte0D1C/UVyb5bImYAFDUpE+DuL+5xuWStiR5b/OUAKC82ld7DFNi3H+ypNdJOsX2NwfHKwrEBYAyBvtZDztqVmI1yL9IqvunBDDRMpizHsFy2zOSNibZuMhp7RWeYATQfaOvBnkkyYbFTmdfUKwBTIS2z1lTrAFMhNrnpIehUwydYvZoiaXtBd4i5eJY2yvpYNJ1XeoUc+yhT86mlx8z9LzVn95Mp5ia1daho1Scq449q3Gc9Xd8Qj9+41GN4xz60e/p8etWN45zwLr79dilBzaOs+y8R3X3GS9qHOfojV+r7vdOp5gF0N0cANqhN+kbOQFA9VL/OuphKNYAOi9iNQgAtELtGzUNQ7EGMBF6jKwBoHLMWQNA/SJWgwBAK3CDEQBqF6vHDUYAqFvEE4wA0ApMgwBA7SJt5wYjANSt3ymGkTUAVI9iDQAtwBOMANACbR9Z0ymmo51ipmT11Px3u9TSfIG3yNIl0vz2AnGmpPle8zh0nBmuS51ijnnq8nxy+oVDzzvxpk10iqlZbR06SsW599UnNI7zzKtv1eOXN+/McsA5j2o+n2ocZ6lfq8e/9Mzm+bzsXn173SmN4zz3uhv1/qPObRznzd/7u+reP91iVoMAQBu0fRqEYg2g8yJuMAJA/cLIGgBawK0fWReZcbe9zvbdtrfavrBETAAoJRo8xTjkqFnjkbXtJZI+LOmlkrZJutX2tUnuahobAEpp+2qQEtmfKGlrknuSPCbpKkmvLBAXAArp7w0y7KhZiWJ9uKT7d3q9bfC1J7C9wfas7dm5ubkClwWA0ST91SDDjpqVKNYL/YT/79GnJDNJppNMr1y5ssBlAWB0bR9Zl1gNsk3S6p1er5L0wwJxAaCYXuU3EIcpUaxvlXSU7SMkPSBpvaQ/KRAXAIoIj5tLSeZtnyfpeklLJF2R5M7GmQFAQbVPcwxT5KGYJJskbSoRCwBKi6QRN2tcbntG0sYkGxczp73FE4wAum/0x80fSbJhsdPZFxRrABOh9qV5w1CsAUyE2h8nH4ZOMR3tFLPE0vYSHV4KdWZZunRK8wUCleo4U6pTTKmOPDXqUqeYZx20Iu8+6rSh5/3xHZ+kU0zNauvQ8aJfeUPjOF/7n4/pqmPPahxn/R2f0Owp6xrHmb7xOv3Xec9qHOdpl35fj32weeeaZW96VNef8OrGcV5+69V0immJto+sKdYAOo/mAwDQBunvD9JmFGsAE4HHzQGgchFPMAJAC1jbKdYAUDduMAJASxR4XGCsKNYAJgJz1gBQuR1tvdqMYg1gInCDEQAqtxf7WVeLYg1gAtTfEHcYijWAicDIGgAqxxOMANASPTZyAoC6Je1fDUKnmMo6xVhWCnQeKdXBpFRHlaWW5ivqXLNE1vaK/p1r1KVOMasOPCznr/qjoeddeM9H6BRTs9o6dHQ1zjvXnN84zrvu+5D+889+vXGcp39kqy4/5k8bxzlny2XV/TvTKWZhzFkDQOVYZw0ALVHJrM4+o1gDmAh0igGAykXS9pbPg0w1+cu232P7O7bvsP0524cUygsACrJ6Ixw1a1SsJd0g6blJjpX0XUkXNU8JAAobdDcfdtSsUbFO8qUk84OXN0ta1TwlAChrx2qQYUfNmo6sd/Z6SV/c3Tdtb7A9a3t2bm6u4GUBYLhehh81G3qD0faXJT1jgW9dnOQLg3MuljQv6crdxUkyI2lG6j/BuE/ZAsA+iNr/uPnQYp3kJXv6vu2zJJ0u6dTU8nwqAOyi7dWp0dI92+skvV3Si5P8okxKAFBe7XPSwzRdZ32ppAMl3TDYEOnmJOc2zgoACurvZz3uLJppVKyTNN9RBwD2g0kfWQNA/Vqw2mMYijWAzuuvBhl3Fs1QrAFMhLbPWdMpprJOMV1VqqNKqY4zXe7wUkotK3FLdIpZecCv5g+fvn7oeX//0AfpFFOz2jp01BbnDSvf1DjOx+Y+2Nl8uhqna5izBoAWaHmtplgD6L6k/ftZU6wBTISW12qKNYDui5izBoBWaHmtplgDmAyMrAGgBSpZNr7PKNYAOi+S5lterSnWACZCu0s1xRrABAi77gFAG0Rp+diaYg2g81hnDQAtwROMAFC5fvOBdg+tKdYAJkIt+3PvK4o1gInQ9mkQOsXQKWaPLFd1F722fLqslpFoiU4xBy85LCc96VVDz7vhFx+lU0zNauvQQRzi1BCna2r6kLd9pKSLJS1PMvxTRNLU4qYEAOMXRfMjHKOwfYXth2x/e5evr7N9t+2tti/cYz7JPUnO2ZufgZE1gO6L1Cs3rfMPki6V9MkdX7C9RNKHJb1U0jZJt9q+VtISSe/e5e+/PslDe3tRijWAiTDiNMgK2zvfUJtJMvOEOMk/2167y987UdLWJPdIku2rJL0yybslnb7vWf8SxRpA50VSb7Ri/fA+3mA8XNL9O73eJum3d3ey7UMl/ZWk42xfNCjqe1SkWNt+q6T3SFqZ5OESMQGgpBGL9b5aaFnZbi+Y5MeSzt2bCzQu1rZXqz9Pc1/TWACwOBZ9I6dtklbv9HqVpB+WvECJ1SDvk/Q2tX+7WAAdFUvz3j70aOBWSUfZPsL2MknrJV1bJPmBRsXa9pmSHkhy+wjnbrA9a3t2bm6uyWUBYK/1lKHHKGx/WtLXJR1te5vtc5LMSzpP0vWStki6OsmdJfMfOg1i+8uSnrHAty6W9A5JLxvlQoM7qjNS/wnGvcgRABrpT4KUeeA8yWt28/VNkjYVucgChhbrJC9Z6Ou2nyfpCEm3Dx7ZXiXpNtsnJnmwaJYA0NAi32BcdPt8gzHJtyQdtuO17X+XNM1qEAA16rndWzmxzhpA5/UnQSjWkqQka0vFAoCyou2aH+XE5bZnJG1MsnGRk9orjKwBdF408jTII0k2LHI6+4RiDWACMA0CAK0QNXroZezoFEOnGGBBXeoUc+CS5Tn8oBcOPe8HP9tEp5ia1dahgzjEqSFOt0Tb9fi4k2iEYg2g8/pbpDJnDQDVK/W4+bhQrAFMgKjX8huMFGsAnRe1f2RNd3MAEyDqZfvQQ4MnGG2fMe6Md8XIGsAEiHqjrQbhCUYAGJdI6qXd0yAUawAToFzzgXGhWAPovkgJq0EAoGrsZw0ALRHmrAGgdlEv7A0CAJXjBiMAVC/iBiMAtECYswaANmj7NMjEd4qZmpqqpiMGUAvb6vXqKG4lOsVM+YAsXXro0PMen/+PrZK+Irqb16eWNySAxbMXc9bsDQIA49PvFdNmFGsAE4EbjABQPdZZA0BLtLtYN+4UY/t823fbvtP235RICgDKipL5oUfNGo2sbf+epFdKOjbJo7YPK5MWAJTW7pF102mQN0r66ySPSlKSh5qnBACLoOXPUzSdBnm2pBfZvsX2V22fsLsTbW+wPWt7dm5uruFlAWBvZKT/ajZ0ZG37y5KescC3Lh78/adJOknSCZKutn1kFngkMMmMpJlBzDnb9zZJfEQrJD28H64zKvLZM/LZs0nN55kFYlwvza8Y4bya/n2foNHj5ravU38a5KbB6+9LOilJFUNn27NNH1MtiXz2jHz2jHwmW9NpkM9LOkWSbD9b0jJV/MkEAG3V9AbjFZKusP1tSY9JOmuhKRAAQDONinWSxyS9tlAui2Fm3Ansgnz2jHz2jHwm2Fi2SAUA7J3GTzACABYfxRoAWmAiinWN+5fYfqvt2B5l7edi5vEe29+xfYftz9k+ZAw5rBv8frbavnB/X3+XXFbb/ortLYP3ywXjzGcH20ts/5vtf6wgl0NsXzN432yx/YJx5zQJOl+sd9m/5Dcl/e2YU5Lt1ZJeKum+ceci6QZJz01yrKTvSrpof17c9hJJH5b0+5KeI+k1tp+zP3PYxbyktyQ5Rv2Hvf58zPnscIGkLeNOYuADkq5L8huSnq968uq0zhdr1bl/yfskvU0a//OtSb6UX243drOkVfs5hRMlbU1yz2B10VXqf7iORZIfJblt8Oefql+IDh9XPpJke5Wk0yRdNs48BrkcLOl3JF0u9VeEJfnvsSY1ISahWI+8f8n+YPtMSQ8kuX2ceezG6yV9cT9f83BJ9+/0epvGXBx3sL1W0nGSbhlzKu9X/8O9hm3jjpQ0J+njg2mZy2w/edxJTYJONB8otX/JfsrnHZJetljX3tt8knxhcM7F6k8BXLk/c5PkBb429v/jsP0USZ+R9OYkPxljHqdLeijJZtu/O648drJU0vGSzk9yi+0PSLpQ0jvHm1b3daJYJ3nJ7r5n+42SPjsozt+w3VN/A5pF279kd/nYfp6kIyTdblvqTzncZvvEJA/u73x2yussSadLOnUMT6Buk7R6p9erJP1wP+fwBLYPUL9QX5nks+PMRdLJks60/QpJT5J0sO1PJRnXw2jbJG1LsuP/Nq5Rv1hjkU3CNMjnVcn+JUm+leSwJGuTrFX/jX/8YhbqYWyvk/R2SWcm+cUYUrhV0lG2j7C9TNJ6SdeOIQ9JkvufopdL2pLkvePKY4ckFyVZNXi/rJd04xgLtQbv1fttHz340qmS7hpXPpOkEyPrIdi/ZM8ulXSgpBsGo/2bk5y7vy6eZN72eZKul7RE0hVJ7txf11/AyZJeJ+lbtr85+No7kmwaX0rVOV/SlYMP13sknT3mfCYCj5sDQAtMwjQIALQexRoAWoBiDQAtQLEGgBagWANAC1CsAaAFKNYA0AL/B2oMpMj2y2tkAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_event(cells[8])\n"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "f5c7603ba5c706ac0ece86f855ab939dae7cd06332146575f448356a26e0dec1"
  },
  "kernelspec": {
   "display_name": "Python 3.9.12 ('ai')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}