{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uproot\n",
    "import awkward as ak\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "import mplhep\n",
    "mplhep.style.use([\"LHCbTex2\"])\n",
    "input_tree = uproot.open({\"/work/guenther/reco_tuner/data/param_data_selected.root\": \"Selected\"})\n",
    "array = input_tree.arrays()\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\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(exptext: Custom Text(0.05, 0.95, 'LHCb'),\n",
       " expsuffix: Custom Text(0.05, 0.955, 'Simulation'))"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 56996 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "plt.hist(array[\"z_mag_x_fringe\"], bins=50,\n",
    "        range=[5150,5300], color='#2A9D8F', density=True)\n",
    "plt.xlabel(r\"z$_{Mag}$ [mm]\")\n",
    "plt.ylabel(\"Number of Tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(exptext: Custom Text(0.05, 0.95, 'LHCb'),\n",
       " expsuffix: Custom Text(0.05, 0.955, 'Simulation'))"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 62600 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace( -0.4, 0.4, 50 )\n",
    "sns.regplot(x=ak.to_numpy(array[\"tx\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)\n",
    "plt.xlabel(\"dx/dz(VELO)\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(exptext: Custom Text(0.05, 0.95, 'LHCb'),\n",
       " expsuffix: Custom Text(0.05, 0.955, 'Simulation'))"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 55820 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace( -0.25, 0.25, 50 )\n",
    "sns.regplot(x=ak.to_numpy(array[\"ty\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)\n",
    "plt.xlabel(\"dy/dz(VELO)\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#bins = np.linspace( -300, 300, 50 )\n",
    "#sns.regplot(x=ak.to_numpy(array[\"x\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#bins = np.linspace( -300, 300, 50 )\n",
    "#sns.regplot(x=ak.to_numpy(array[\"y\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#bins = np.linspace( -1.0, 1.0, 50 )\n",
    "#sns.regplot(x=ak.to_numpy(array[\"dSlope_out\"]), y=ak.to_numpy(array[\"z_mag_x\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(exptext: Custom Text(0.05, 0.95, 'LHCb'),\n",
       " expsuffix: Custom Text(0.05, 0.955, 'Simulation'))"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/work/guenther/reco_tuner/env/tuner_env/envs/tuner/lib/python3.10/site-packages/IPython/core/pylabtools.py:151: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.\n",
      "  fig.canvas.print_figure(bytes_io, **kw)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 58756 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace( -1, 1, 50 )\n",
    "sns.regplot(x=ak.to_numpy(array[\"dSlope_fringe\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)\n",
    "plt.xlabel(\"$\\Delta$dx/dz\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(exptext: Custom Text(0.05, 0.95, 'LHCb'),\n",
       " expsuffix: Custom Text(0.05, 0.955, 'Simulation'))"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/work/guenther/reco_tuner/env/tuner_env/envs/tuner/lib/python3.10/site-packages/IPython/core/pylabtools.py:151: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.\n",
      "  fig.canvas.print_figure(bytes_io, **kw)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 82116 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#import matplotlib.pyplot as plt\n",
    "bins = np.linspace( -2000, 2000, 50 )\n",
    "sns.regplot(x=ak.to_numpy(array[\"x_l0\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=50, fit_reg=None, x_estimator=np.mean, label=\"T1X1\")\n",
    "sns.regplot(x=ak.to_numpy(array[\"x_l4\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=50, fit_reg=None, x_estimator=np.mean, label=\"T2X1\")\n",
    "sns.regplot(x=ak.to_numpy(array[\"x_l8\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=50, fit_reg=None, x_estimator=np.mean, label=\"T3X1\")\n",
    "plt.legend()\n",
    "plt.xlabel(\"x [mm]\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['tx^2' 'tx dSlope_fringe' 'ty^2' 'dSlope_fringe^2']\n",
      "intercept= 5205.144186525624\n",
      "coef= {'tx^2': -320.7206595710594, 'tx dSlope_fringe': 702.1384894815535, 'ty^2': -316.36350963107543, 'dSlope_fringe^2': 441.59909857558097}\n",
      "r2 score= 0.9604900589467942\n",
      "RMSE = 8.772908410819978\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",
    "import numpy as np\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"dSlope_fringe\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\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=2, 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",
    "keep = [\n",
    "    #\"tx\",\n",
    "    #\"ty\",\n",
    "    #\"dSlope_fringe\",\n",
    "    \"tx^2\",\n",
    "    \"tx dSlope_fringe\",\n",
    "    \"ty^2\",\n",
    "    \"dSlope_fringe^2\"\n",
    "]\n",
    "remove = [i for i, f in enumerate(poly_features) if f not in keep]\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",
    "\n",
    "lin_reg = LinearRegression()#Lasso(alpha=0.01)\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(exptext: Custom Text(0.0, 1, 'LHCb'),\n",
       " expsuffix: Custom Text(0.0, 1.005, 'Simulation'))"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 73664 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace( 5150,5300, 30 )\n",
    "ax = sns.regplot(x=y_test, y=abs(y_test-y_pred_test), x_bins=bins, fit_reg=None, x_estimator=np.mean, label=\"bla\")\n",
    "ax2 = ax.twinx()\n",
    "ax2.hist(y_test, bins=30,\n",
    "        range=[5150,5300], color='#2A9D8F', alpha=0.8, align='left')\n",
    "ax.set_xlabel(r\"z$_{Mag}$ [mm]\")\n",
    "ax.set_ylabel(\"Mean Deviation [mm]\")\n",
    "ax2.set_ylabel(\"Number of Tracks\")\n",
    "mplhep.lhcb.text(\"Simulation\", loc=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "median_z_mag_x = np.median(array[\"z_mag_x_fringe\"])\n",
    "print(median_z_mag_x)\n",
    "params_per_layer = [[] for _ in range(12)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_array(name, intercept, coef):\n",
    "    coef = [str(c)+\"f\" for c in coef if c != 0.0]\n",
    "    intercept = str(intercept) + \"f\"\n",
    "    code = f\"constexpr std::array {name}\"\n",
    "    code += \"{\" + \", \".join([intercept]+list(coef)) +\"};\"\n",
    "    return code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "array[\"x_diff_straight_l0\"] = array[\"x_l0\"] - array[\"x\"] - array[\"tx\"] * ( array[\"z_l0\"] - array[\"z\"])\n",
    "array[\"x_l0_rel\"] = array[\"x_l0\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    #\"x_l0_rel\",\n",
    "    \"x_diff_straight_l0\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\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=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "poly_features = poly.get_feature_names_out(input_features=features)\n",
    "keep = [\n",
    "    #\"tx\",\n",
    "    #\"ty\",\n",
    "    #\"x_l0_rel\",\n",
    "    \"tx^2\",\n",
    "    #\"tx x_l0_rel\",\n",
    "    \"tx x_diff_straight_l0\",\n",
    "    \"ty^2\",\n",
    "    #\"x_l0_rel^2\"\n",
    "    \"x_diff_straight_l0^2\"\n",
    "]\n",
    "remove = [i for i, f in enumerate(poly_features) if f not in keep]\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()#Lasso(alpha=0.004)\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(\"zMagnetParams_l0\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[0] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l1_rel\"] = array[\"x_l1\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l1_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\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=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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(\"zMagnetParams_l1\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[1] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l2_rel\"] = array[\"x_l2\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l2_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\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=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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(\"zMagnetParams_l2\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[2] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l3_rel\"] = array[\"x_l3\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l3_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\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=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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(\"zMagnetParams_l3\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[3] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l4_rel\"] = array[\"x_l4\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l4_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\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=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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(\"zMagnetParams_l4\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[4] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l5_rel\"] = array[\"x_l5\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l5_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\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=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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(\"zMagnetParams_l5\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[5] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l6_rel\"] = array[\"x_l6\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l6_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\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=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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(\"zMagnetParams_l6\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[6] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l7_rel\"] = array[\"x_l7\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l7_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\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=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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(\"zMagnetParams_l7\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[7] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l8_rel\"] = array[\"x_l8\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l8_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\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=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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(\"zMagnetParams_l8\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[8] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l9_rel\"] = array[\"x_l9\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l9_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\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=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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(\"zMagnetParams_l9\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[9] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l10_rel\"] = array[\"x_l10\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l10_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\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=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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(\"zMagnetParams_l10\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[10] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l11_rel\"] = array[\"x_l11\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l11_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\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=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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(\"zMagnetParams_l11\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[11] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"dSlope_fringe\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\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=2, 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",
    "keep = [\n",
    "    #\"tx\",\n",
    "    #\"ty\",\n",
    "    #\"dSlope_fringe\",\n",
    "    \"tx^2\",\n",
    "    \"tx dSlope_fringe\",\n",
    "    \"ty^2\",\n",
    "    \"dSlope_fringe^2\"\n",
    "]\n",
    "remove = [i for i, f in enumerate(poly_features) if f not in keep]\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",
    "\n",
    "lin_reg = LinearRegression()#Lasso(alpha=0.01)\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(\"zMagnetParams_dSlope\", lin_reg.intercept_, lin_reg.coef_))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.optimize\n",
    "def parabola(x, a,b,c):\n",
    "     return a*x**2 + b * x + c\n",
    "params_1 = np.array([p[1] / params_per_layer[0][1] for p in params_per_layer])\n",
    "x = [array[f\"z_l{n}\"][0] - array[\"z_ref\"][0] for n in range(12)]\n",
    "print(params_1)\n",
    "print(x)\n",
    "plt.plot(x, params_1, 'o')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params_3 = np.array([params_per_layer[0][3] / p[3] for p in params_per_layer])\n",
    "x = [array[f\"z_l{n}\"][0] - array[\"z_ref\"][0] for n in range(12)]\n",
    "print(params_3**2)\n",
    "print(x)\n",
    "plt.plot(x, params_3, 'o')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.preprocessing import PolynomialFeatures\n",
    "from sklearn.linear_model import LinearRegression, Lasso\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import mean_squared_error\n",
    "feautures = [\"tx\", \"ty\", \"dSlope\"]\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(\n",
    "    data,\n",
    "    target,\n",
    "    test_size=0.2,\n",
    "    random_state=42,\n",
    ")\n",
    "poly = PolynomialFeatures(degree=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform(X_train)\n",
    "X_test_model = poly.fit_transform(X_test)\n",
    "lin_reg = LinearRegression() # or Lasso if regularisation is needed\n",
    "lin_reg.fit(X_train_model, y_train)\n",
    "y_pred_test = lin_reg.predict(X_test_model)\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))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6",
   "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
}
