{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uproot\n",
    "import awkward as ak\n",
    "input_tree_md = uproot.open({\"/work/guenther/reco_tuner/data/param_data_MD_selected_8520.root\": \"Selected\"})\n",
    "input_tree_mu = uproot.open({\"/work/guenther/reco_tuner/data/param_data_MU_selected_8520.root\": \"Selected\"})\n",
    "# this is an event list of dictionaries containing awkward arrays\n",
    "array_md = input_tree_md.arrays()\n",
    "array_mu = input_tree_mu.arrays()\n",
    "array = ak.concatenate([array_md, array_mu])\n",
    "array[\"dSlope_fringe\"] = array[\"tx_ref\"] - array[\"tx\"]\n",
    "array[\"z_mag_x_fringe\"] = (array[\"x\"] - array[\"x_ref\"] - array[\"tx\"] * array[\"z\"] + array[\"tx_ref\"] * array[\"z_ref\"] ) / array[\"dSlope_fringe\"]\n",
    "def format_array(name, coef):\n",
    "    coef = [str(c)+\"f\" for c in coef if c != 0.0]\n",
    "    code = f\"constexpr std::array {name}\"\n",
    "    code += \"{\" + \", \".join(list(coef)) +\"};\"\n",
    "    return code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<AxesSubplot: >"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 16992 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import seaborn as sns\n",
    "import numpy as np\n",
    "bins = 50#np.linspace( -1.5, 1.5, 50 )\n",
    "sns.regplot(x=ak.to_numpy(array[\"dSlope_fringe\"]), y=ak.to_numpy(array[\"CX_ex\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<AxesSubplot: >"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 15724 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = 50#np.linspace( -1., 1., 50 )\n",
    "sns.regplot(x=ak.to_numpy(array[\"tx\"]), y=ak.to_numpy(array[\"CX_ex\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<AxesSubplot: >"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 15664 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = 50#np.linspace( -1., 1., 50 )\n",
    "sns.regplot(x=ak.to_numpy(array[\"x\"]), y=ak.to_numpy(array[\"CX\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<AxesSubplot: >"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 15784 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = 50#np.linspace( -1., 1., 50 )\n",
    "sns.regplot(x=ak.to_numpy(array[\"dSlope_fringe\"]), y=ak.to_numpy(array[\"DX\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<AxesSubplot: >"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 18472 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = 50#np.linspace( -1., 1., 50 )\n",
    "sns.regplot(x=ak.to_numpy(array[\"tx\"]), y=ak.to_numpy(array[\"DX\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<AxesSubplot: >"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 18060 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = 50#np.linspace( -1., 1., 50 )\n",
    "sns.regplot(x=ak.to_numpy(array[\"x\"]), y=ak.to_numpy(array[\"DX\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['dSlope_fringe' 'tx dSlope_fringe' 'ty dSlope_fringe'\n",
      " 'tx^2 dSlope_fringe' 'tx ty dSlope_fringe' 'ty^2 dSlope_fringe']\n",
      "intercept= 0.0\n",
      "coef= {'dSlope_fringe': 2.3340814100448438e-05, 'tx dSlope_fringe': 8.310801914170572e-09, 'ty dSlope_fringe': -1.025087641029144e-06, 'tx^2 dSlope_fringe': 8.228862899700043e-06, 'tx ty dSlope_fringe': -1.5662848683486495e-06, 'ty^2 dSlope_fringe': -0.0003728935994887431}\n",
      "r2 score= 0.9860788549938552\n",
      "RMSE = 6.07178841864582e-07\n",
      "constexpr std::array cx_params{2.3340814100448438e-05f, 8.310801914170572e-09f, -1.025087641029144e-06f, 8.228862899700043e-06f, -1.5662848683486495e-06f, -0.0003728935994887431f};\n"
     ]
    }
   ],
   "source": [
    "from sklearn.preprocessing import PolynomialFeatures\n",
    "from sklearn.linear_model import LinearRegression, Lasso, Ridge\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.pipeline import Pipeline\n",
    "from sklearn.metrics import mean_squared_error\n",
    "\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\",\n",
    "    \"dSlope_fringe\",\n",
    "]\n",
    "target_feat = \"CX_ex\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=3, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "poly_features = poly.get_feature_names_out(input_features=features)\n",
    "remove = [i for i, f in enumerate(poly_features) if (\"dSlope_fringe\" not in f ) or (\"dSlope_fringe^\" in f)]\n",
    "X_train_model = np.delete( X_train_model, remove, axis=1)\n",
    "X_test_model = np.delete( X_test_model, remove, axis=1)\n",
    "poly_features = np.delete(poly_features, remove )\n",
    "print(poly_features)\n",
    "\n",
    "lin_reg = LinearRegression(fit_intercept=False)\n",
    "lin_reg.fit( X_train_model, y_train)\n",
    "y_pred_test = lin_reg.predict( X_test_model )\n",
    "print(\"intercept=\", lin_reg.intercept_)\n",
    "print(\"coef=\", dict(zip(poly_features, lin_reg.coef_)))\n",
    "print(\"r2 score=\", lin_reg.score(X_test_model, y_test))\n",
    "print(\"RMSE =\", mean_squared_error(y_test, y_pred_test, squared=False))\n",
    "print(format_array(\"cx_params\", lin_reg.coef_))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['dSlope_fringe' 'tx dSlope_fringe' 'ty dSlope_fringe'\n",
      " 'tx^2 dSlope_fringe' 'tx ty dSlope_fringe' 'ty^2 dSlope_fringe']\n",
      "intercept= 0.0\n",
      "coef= {'dSlope_fringe': -7.054315710920644e-09, 'tx dSlope_fringe': 2.22196526392023e-11, 'ty dSlope_fringe': 6.104698063377923e-10, 'tx^2 dSlope_fringe': 2.85681046841909e-09, 'tx ty dSlope_fringe': 3.789175266592706e-10, 'ty^2 dSlope_fringe': 9.907774160587578e-08}\n",
      "r2 score= 0.97610245571401\n",
      "RMSE = 2.4251033068807623e-10\n",
      "constexpr std::array dx_params{-7.054315710920644e-09f, 2.22196526392023e-11f, 6.104698063377923e-10f, 2.85681046841909e-09f, 3.789175266592706e-10f, 9.907774160587578e-08f};\n"
     ]
    }
   ],
   "source": [
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\",\n",
    "    \"dSlope_fringe\",\n",
    "]\n",
    "target_feat = \"DX_ex\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=3, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "poly_features = poly.get_feature_names_out(input_features=features)\n",
    "remove = [i for i, f in enumerate(poly_features) if (\"dSlope_fringe\" not in f ) or (\"dSlope_fringe^\" in f)]\n",
    "X_train_model = np.delete( X_train_model, remove, axis=1)\n",
    "X_test_model = np.delete( X_test_model, remove, axis=1)\n",
    "poly_features = np.delete(poly_features, remove )\n",
    "print(poly_features)\n",
    "\n",
    "lin_reg = LinearRegression(fit_intercept=False)\n",
    "lin_reg.fit( X_train_model, y_train)\n",
    "y_pred_test = lin_reg.predict( X_test_model )\n",
    "print(\"intercept=\", lin_reg.intercept_)\n",
    "print(\"coef=\", dict(zip(poly_features,lin_reg.coef_)))\n",
    "print(\"r2 score=\", lin_reg.score(X_test_model, y_test))\n",
    "print(\"RMSE =\", mean_squared_error(y_test, y_pred_test, squared=False))\n",
    "print(format_array(\"dx_params\", lin_reg.coef_))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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
}
