{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uproot\n",
    "import awkward as ak\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "input_tree = uproot.open({\"/work/guenther/reco_tuner/data/param_data_selected_all_p.root\": \"Selected\"})\n",
    "array = input_tree.arrays()\n",
    "array = array[[field for field in ak.fields(array) if \"scifi_hit\" not in field]]\n",
    "df = ak.to_pandas(array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 39844 characters]",
      "text/plain": [
       "<Figure size 800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "urax = df.loc[(df[\"x_ref\"] > -3200) & (df[\"x_ref\"] < 3200), \"x_ref\"].plot.hist(bins=100, logy=True,title='x positions on reference plane', figsize=(8,6), fontsize=20,edgecolor='black')\n",
    "urax.set_xlabel('x [mm]', fontsize=24)\n",
    "urax.set_ylabel('Count', fontsize=24)\n",
    "urax.title.set_size(24)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import curve_fit\n",
    "import numpy as np\n",
    "def fitFastSigmoid(x, p0, p1, p2 ):\n",
    "    return p0 + p1 * x / (1 + abs(x * p2) ) \n",
    "\n",
    "def sqrtSigmoid( x, p0, p1, p2 ):\n",
    "    return p0 + p1 * x / np.sqrt( p2 + x**2 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-1152.3191321144784 1144.5112688972854\n",
      "BINS=900\n",
      "shift=126.0\n",
      "[ 4.50871394e+02  5.78026621e-01 -6.72848459e-04]\n",
      "at -3000 = 2\n",
      "at  0 = 576\n",
      "at 3000 = 1151\n",
      "at -3000 = 2.3960602381498575\n",
      "at  0 = 576.8713937732111\n",
      "at 3000 = 1151.3467273082724\n",
      "\n",
      "-1152.3191321144784 1145.2265343593158\n",
      "BINS=1000\n",
      "shift=140.0\n",
      "[5.00965281e+02 6.42244329e-01 6.72824292e-04]\n",
      "at -3000 = 2\n",
      "at  0 = 640\n",
      "at 3000 = 1279\n",
      "at -3000 = 2.6514499904290005\n",
      "at  0 = 640.965281389201\n",
      "at 3000 = 1279.279112787973\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>NBINS</th>\n",
       "      <th>nBinsReal</th>\n",
       "      <th>minBinWidth</th>\n",
       "      <th>maxBinWidth</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>900</td>\n",
       "      <td>1152</td>\n",
       "      <td>1.520524</td>\n",
       "      <td>411.883522</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1000</td>\n",
       "      <td>1280</td>\n",
       "      <td>1.378500</td>\n",
       "      <td>406.784548</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   NBINS  nBinsReal  minBinWidth  maxBinWidth\n",
       "0    900       1152     1.520524   411.883522\n",
       "1   1000       1280     1.378500   406.784548"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 135820 characters]",
      "text/plain": [
       "<Figure size 2400x2400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from math import ceil\n",
    "selection = (df[\"x_ref\"] > -3000) & (df[\"x_ref\"] < 3000) \n",
    "data = df.loc[selection, \"x_ref\"]\n",
    "#func = sqrtSigmoid\n",
    "func = fitFastSigmoid\n",
    "#NBINS=[800, 900, 1000, 1100]\n",
    "NBINS=[900, 1000]\n",
    "#SHIFTS = [48, 50, 53, 52]\n",
    "SHIFTS=[126., 140.]\n",
    "minBinWidths = []\n",
    "maxBinWidths = []\n",
    "nBinsReal = []\n",
    "fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(24, 24))\n",
    "for i, BINS in enumerate(NBINS):\n",
    "    out, equal_bins = pd.qcut(data, q=BINS, retbins=True)\n",
    "    binWidths = np.array([ abs(equal_bins[j] - equal_bins[j+1]) for j in range(len(equal_bins)-1)])\n",
    "    minBinWidths.append( binWidths.min() )\n",
    "    maxBinWidths.append( binWidths.max() )\n",
    "    bin_numbering = np.arange(0, BINS+1)\n",
    "    equalbins_center = equal_bins[int(BINS/10):int(9*BINS/10)]\n",
    "    bin_numbering_center = bin_numbering[int(BINS/10):int(9*BINS/10)]\n",
    "    print(equalbins_center[0], equalbins_center[-1])\n",
    "    popt, pcov = curve_fit( func, xdata=equalbins_center, ydata=bin_numbering_center )\n",
    "    if i<2:\n",
    "        axes[0,i].plot(equal_bins, bin_numbering, '.', label='data')\n",
    "        axes[0,i].plot(equal_bins, func(equal_bins, *popt), label=f'p0={popt[0]}\\np1={popt[1]}\\np2={popt[2]}\\n')\n",
    "        axes[0,i].legend()\n",
    "    else:\n",
    "        axes[1,i-2].plot(equal_bins, bin_numbering, '.', label='data')\n",
    "        axes[1,i-2].plot(equal_bins, func(equal_bins, *popt), label=f'p0={popt[0]}\\np1={popt[1]}\\np2={popt[2]}\\n')\n",
    "        axes[1,i-2].legend()\n",
    "    shift = SHIFTS[i]\n",
    "    nBinsReal.append(ceil(func(3000, popt[0]+shift, popt[1], popt[2])))\n",
    "    print(f'BINS={BINS}')\n",
    "    print(f'shift={shift}')\n",
    "    print(popt)\n",
    "    print(f'at -3000 = {int(func(-3000., popt[0]+shift, popt[1], popt[2]))}')\n",
    "    print(f'at  0 = {int(func(0., popt[0]+shift, popt[1], popt[2]))}')\n",
    "    print(f'at 3000 = {int(func(3000., popt[0]+shift, popt[1], popt[2]))}')\n",
    "    print(f'at -3000 = {func(-3000., popt[0]+shift, popt[1], popt[2])}')\n",
    "    print(f'at  0 = {func(0., popt[0]+shift, popt[1], popt[2])}')\n",
    "    print(f'at 3000 = {func(3000., popt[0]+shift, popt[1], popt[2])}')\n",
    "    print()\n",
    "binData = pd.DataFrame({'NBINS':NBINS, 'nBinsReal': nBinsReal,'minBinWidth':minBinWidths, 'maxBinWidth':maxBinWidths})\n",
    "#ax = binData.plot(x='NBINS', y=['minBinWidth', 'maxBinWidth'], logy=True)\n",
    "binData"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 (conda)",
   "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.10.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "a2eff8b4da8b8eebf5ee2e5f811f31a557e0a202b4d2f04f849b065340a6eda6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
