{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d52f88aa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Read and plot the simulation results of particle interactions in Oriented Crystals\n",
    "# obatined through example ch2, which is baed on G4ChannelingFastSimModel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c010d890",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import uproot\n",
    "from matplotlib.colors import LogNorm  # optional, for log color scaling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bf061146",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#################################### INPUT FILE #########################################\n",
    "# Set path and filename of the simulation file\n",
    "G4_sim_path = \"\"\n",
    "root_file = \"results\"\n",
    "\n",
    "# Set whether to save plots (without displaying them) or just display them\n",
    "save_fig = True\n",
    "fig_path = G4_sim_path\n",
    "#########################################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3b298eab",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rf_content: ['crystal', 'detector_primaries', 'detector_photons', 'detector_secondaries', 'missed_crystal'] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Create directory where to strore the figures if it does not exist\n",
    "if fig_path != '' and not os.path.exists(fig_path):\n",
    "    os.makedirs(fig_path)\n",
    "    print('created fig_path:', fig_path)\n",
    "    \n",
    "# Open the simulation output root file    \n",
    "rf = uproot.open(G4_sim_path + root_file + '.root')\n",
    "rf_content = [item.split(';')[0] for item in rf.keys()]\n",
    "print('rf_content:', rf_content, '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "599eb756",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Import the scoring ntuples and convert them into pandas dataframes\n",
    "branches = [\"eventID\", \"volume\", \"x\", \"y\", \"angle_x\", \"angle_y\", \\\n",
    "            \"Ekin\" , \"particle\", \"particleID\", \"parentID\"]\n",
    "branchesprimary = branches + [\"incoming_angle_x\", \"deflection_angle_x\", \\\n",
    "                              \"incoming_angle_y\", \"deflection_angle_y\"]\n",
    "\n",
    "df_in = rf['crystal'].arrays(branches, library='pd')\n",
    "df_prim = rf['detector_primaries'].arrays(branchesprimary, library='pd')\n",
    "df_ph = rf['detector_photons'].arrays(branches, library='pd')\n",
    "df_sec = rf['detector_secondaries'].arrays(branches, library='pd')\n",
    "df_missed = rf['missed_crystal'].arrays(branches, library='pd')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "07c70dd2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#########################################################################################\n",
    "# Plot angle_x distribution of primaries at the detector after interaction with a crystal\n",
    "\n",
    "############# INPUT #############\n",
    "# Feel free to modify according to your needs\n",
    "Nmax = 100000000 #max number of events to elaborate\n",
    "\n",
    "# Feel free to replace df_prim by df_in, df_ph, df_sec or df_missed\n",
    "# Feel free to replace \"angle_x\" by other ntuples from branches and\n",
    "# from branchesprimary (for df_prim) \n",
    "# ONLY NUMERIC VALUES\n",
    "datax = df_prim[\"angle_x\"][:Nmax]*1.e3  #mrad <= rad (feel free to modify the coefficient)\n",
    "\n",
    "# Feel free to modify the number of bins and the plot range\n",
    "NbinTheta = 100\n",
    "rangeTheta = [-1, 2] #mrad\n",
    "\n",
    "# Set whether to use linear o log scale\n",
    "use_log_y = False  # set True for LogNorm color scale\n",
    "\n",
    "# Feel free to modify the names of axes\n",
    "plt_xlabel = '$\\\\theta_x$ [mrad]'\n",
    "plt_ylabel = 'PDF:  1/N dN/d$\\\\theta_x$ [mrad]$^{-1}$'\n",
    "\n",
    "# Feel free to modify the filename to save the plot\n",
    "filename = 'thetaXdistribution.pdf'\n",
    "\n",
    "#some plt parameters\n",
    "fs = 16\n",
    "lw = 2\n",
    "#################################\n",
    "\n",
    "# Create 1D histogram\n",
    "thetaXdistrib, thetaEdges = np.histogram(datax.values, \\\n",
    "                                         bins=NbinTheta, range=rangeTheta, density=True)\n",
    "thetabin = thetaEdges[:-1] + (thetaEdges[1]-thetaEdges[0])*0.5\n",
    "plt.figure(figsize=(9, 6))\n",
    "plt.grid()\n",
    "plt.plot(thetabin, thetaXdistrib, linewidth=lw, alpha=1, label='')\n",
    "plt.xlabel(plt_xlabel, fontsize=fs)\n",
    "plt.ylabel(plt_ylabel, fontsize=fs)\n",
    "\n",
    "# Set log scale\n",
    "if use_log_y:\n",
    "    plt.yscale('log',base=2) \n",
    "\n",
    "# Save the plot or just show it\n",
    "if save_fig:\n",
    "    plt.savefig(fig_path + filename)\n",
    "    plt.close() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c9f9c18a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#########################################################################################\n",
    "# angle_x_in - angle_x_defl distribution of primaries at the detector after interaction with a crystal\n",
    "\n",
    "############# INPUT #############\n",
    "# Feel free to modify according to your needs\n",
    "Nmax = 100000000 #max number of events to elaborate\n",
    "\n",
    "# Example data (replace these with your real arrays)\n",
    "# datatetaxin and datatetadeflx must be the same length\n",
    "datatetaxin = df_prim[\"incoming_angle_x\"][:Nmax]*1.e3  #mrad <= rad (feel free to modify the coefficient)\n",
    "datatetadeflx = df_prim[\"deflection_angle_x\"][:Nmax]*1.e3  #mrad <= rad (feel free to modify the coefficient)\n",
    "\n",
    "# Feel free to modify the number of bins and the plot range\n",
    "NbinTheta = 50\n",
    "\n",
    "# Feel free to modify the plot range\n",
    "xrange = (-0.1, 0.1)\n",
    "yrange = (-1, 2)\n",
    "\n",
    "# Set whether to use linear o log scale\n",
    "use_log_color = True  # set True for LogNorm color scale\n",
    "\n",
    "# Feel free to modify the names of axes\n",
    "plt_xlabel2 = '$\\\\theta_{x in}$ [mrad]'\n",
    "plt_ylabel2 = '$\\\\theta_{x defl}$ [mrad]'\n",
    "\n",
    "# Feel free to modify the filename to save the plot\n",
    "filename2 = 'thetaXin_thetaXdefl.pdf'\n",
    "\n",
    "#some plt parameters\n",
    "fs = 16\n",
    "lw = 2\n",
    "#################################\n",
    "\n",
    "# Create 2D histogram\n",
    "plt.figure(figsize=(8, 6))\n",
    "hist = plt.hist2d(\n",
    "    datatetaxin,\n",
    "    datatetadeflx,\n",
    "    bins=NbinTheta,\n",
    "    density=True,\n",
    "    range=[xrange, yrange],\n",
    "    norm=LogNorm() if use_log_color else None,\n",
    "    cmap='jet'\n",
    ")\n",
    "\n",
    "# Add colorbar (PDF scale)\n",
    "cbar = plt.colorbar()\n",
    "cbar.set_label('PDF', fontsize=fs)\n",
    "\n",
    "# Labels and title\n",
    "plt.xlabel(plt_xlabel2, fontsize=fs)\n",
    "plt.ylabel(plt_ylabel2, fontsize=fs)\n",
    "\n",
    "# Save the plot or just show it\n",
    "if save_fig:\n",
    "    plt.savefig(fig_path + filename2)\n",
    "    plt.close() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3098a395",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "################################################################################################\n",
    "# Plot spectrum\n",
    "\n",
    "############# INPUT #############\n",
    "# Feel free to modify the collimator parameters\n",
    "# !!! related only to real secondary photons from results.root, not from Spectrum.dat\n",
    "# For Spectrum.dat, see the options in the simulation macro.\n",
    "apply_collimation = True\n",
    "coll_angle = 2.3183 #mrad\n",
    "\n",
    "# Feel free to modify\n",
    "NbinE = 20\n",
    "rangeE = [0, 10] #MeV\n",
    "\n",
    "# path of the spectrum file obtained using all the Baier-Katkov integration photons\n",
    "BK_spectrum_file = \"Spectrum.dat\"\n",
    "#################################\n",
    "\n",
    "# Array with photon energies and angles\n",
    "Eph = df_ph['Ekin'].values #MeV \n",
    "\n",
    "Nph = len(Eph)\n",
    "print(\"number of emitted photons:\", Nph)\n",
    "thetaX_ph = df_ph['angle_x'].values*1e3 #rad -> mrad\n",
    "thetaY_ph = df_ph['angle_y'].values*1e3 #rad -> mrad\n",
    "\n",
    "# Take only the photons inside the collimator acceptance\n",
    "theta_ph = np.sqrt(thetaX_ph**2 + thetaY_ph**2)   \n",
    "if apply_collimation: \n",
    "    thetaX_ph = thetaX_ph[theta_ph <= coll_angle]\n",
    "    thetaY_ph = thetaY_ph[theta_ph <= coll_angle]\n",
    "    Eph = Eph[theta_ph <= coll_angle]\n",
    "    theta_ph = theta_ph[theta_ph <= coll_angle]\n",
    "\n",
    "# Calculate the scored photon energy spectrum\n",
    "spectrum0, EbinEdges = np.histogram(Eph, bins=NbinE, range=rangeE, density=False)\n",
    "Ebin = EbinEdges[:-1] + (EbinEdges[1]-EbinEdges[0])*0.5\n",
    "stepx = Ebin[1]-Ebin[0]\n",
    "Nprimaries = df_in[\"Ekin\"].size\n",
    "\n",
    "spectrum = spectrum0 / (Nprimaries*stepx)\n",
    "spectral_intensity = Ebin * spectrum\n",
    "\n",
    "# Statistical uncertainties: sqrt(N)\n",
    "spectrum_err = np.sqrt(spectrum0) / (Nprimaries * stepx)\n",
    "spectral_intensity_err = Ebin * spectrum_err\n",
    "\n",
    "# Read the spectrum file obtained using all the Bair-Katkov integration photons\n",
    "BK_spectrum = np.loadtxt(G4_sim_path+BK_spectrum_file, dtype='float', comments='#', \\\n",
    "                         delimiter=' ', skiprows=1, unpack=True)\n",
    "E_ext = BK_spectrum[0]\n",
    "S_ext = BK_spectrum[1]\n",
    "\n",
    "# Plot the photon energy spectrum\n",
    "fig = plt.figure(figsize=(13, 6))\n",
    "fs = 16\n",
    "lw = 2\n",
    "bw = 0.6\n",
    "color0 = '#B1B3FB'\n",
    "\n",
    "#!!! The spectrum is normalized on the total radiation probability W_rad \n",
    "#(from MinPhotonEnergy to the energy of the primary particle) which is equivalent to the radiation yield.\n",
    "#The integral is equal to W_rad, not to 1!\n",
    "plt.subplot(1,2,1)\n",
    "plt.bar(Ebin, spectrum, width=bw, color=color0, linewidth=lw, alpha=1, label='secondary photons')\n",
    "plt.errorbar(Ebin, spectrum, yerr=spectrum_err, fmt='o', color='k', capsize=3)\n",
    "plt.xlim(rangeE)\n",
    "plt.plot(E_ext, S_ext, 'r-', lw=2.5, label='from '+BK_spectrum_file)\n",
    "plt.title('Emitted photon spectrum')\n",
    "plt.xlabel('$E$ [MeV]', fontsize=fs)\n",
    "plt.ylabel('$dW_{rad}/dE$ [MeV$^{-1}$]', fontsize=fs)\n",
    "plt.legend()\n",
    "#plt.yscale('log')\n",
    "\n",
    "#The spectral intensity is the spectrum above multiplied by the energy,\n",
    "#so the spectral intensity of bremsstrahlung is nearly constant.\n",
    "plt.subplot(1,2,2)\n",
    "plt.bar(Ebin, spectral_intensity, width=bw, color=color0, linewidth=lw, alpha=1, label='secondary photons')\n",
    "plt.errorbar(Ebin, spectral_intensity, yerr=spectral_intensity_err, fmt='o', color='k', capsize=3)\n",
    "plt.plot(E_ext, E_ext * S_ext, 'r-', lw=2.5, label='from '+BK_spectrum_file)\n",
    "plt.title('Emitted photon spectral intensity')\n",
    "plt.xlabel('$E$ [MeV]', fontsize=fs)\n",
    "plt.ylabel('$E dW_{rad}/dE$', fontsize=fs)\n",
    "plt.xlim(rangeE)\n",
    "plt.legend()\n",
    "#plt.yscale('log')\n",
    "if save_fig:\n",
    "    plt.savefig(fig_path + 'spectrum.pdf')\n",
    "    plt.close() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ddc1e1f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}