{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 106,
   "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",
    "\n",
    "from sklearn.preprocessing import PolynomialFeatures\n",
    "from sklearn.linear_model import LinearRegression, Lasso, Ridge, ElasticNet\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",
    "mplhep.style.use([\"LHCbTex2\"])\n",
    "input_tree = uproot.open(\n",
    "    {\n",
    "        \"/work/cetin/LHCb/reco_tuner/data/tracking_losses_ntuple_B_BJpsi_def_selected.root\": \"Selected\"\n",
    "    }\n",
    ")\n",
    "array = input_tree.arrays()\n",
    "\n",
    "array[\"yStraightOut\"] = array[\"ideal_state_770_y\"] + array[\"ideal_state_770_ty\"] * (\n",
    "    array[\"ideal_state_10000_z\"] - array[\"ideal_state_770_z\"]\n",
    ")\n",
    "array[\"yDiffOut\"] = array[\"ideal_state_10000_y\"] - array[\"yStraightOut\"]\n",
    "array[\"yStraightEndT\"] = array[\"ideal_state_770_y\"] + array[\"ideal_state_770_ty\"] * (\n",
    "    9410.0 - array[\"ideal_state_770_z\"]\n",
    ")\n",
    "array[\"yDiffEndT\"] = array[\"ideal_state_9410_y\"] - array[\"yStraightEndT\"]\n",
    "\n",
    "array[\"dSlope_xEndT\"] = array[\"ideal_state_9410_tx\"] - array[\"ideal_state_770_tx\"]\n",
    "array[\"dSlope_yEndT\"] = array[\"ideal_state_9410_ty\"] - array[\"ideal_state_770_ty\"]\n",
    "array[\"dSlope_xEndT_abs\"] = abs(array[\"dSlope_xEndT\"])\n",
    "array[\"dSlope_yEndT_abs\"] = abs(array[\"dSlope_yEndT\"])\n",
    "\n",
    "sel_array = array[(array[\"yDiffOut\"] < 100) & (array[\"yDiffOut\"] > -100)]\n",
    "# sel_array = array\n",
    "\n",
    "\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": 107,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 56332 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "plt.hist(\n",
    "    sel_array[\"yDiffOut\"],\n",
    "    bins=100,\n",
    "    # range=[5100, 5700],\n",
    "    color=\"#2A9D8F\",\n",
    "    density=True,\n",
    ")\n",
    "plt.xlabel(r\"$y_{corr}$ [mm]\")\n",
    "plt.ylabel(\"Number of Tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_dist.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 53184 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-200, 200, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_770_y\"]),\n",
    "    y=ak.to_numpy(sel_array[\"yDiffOut\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"y (VELO)\")\n",
    "plt.ylabel(\"$y_{corr}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 49448 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-200, 200, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_9410_y\"]),\n",
    "    y=ak.to_numpy(sel_array[\"yDiffOut\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"y (T)\")\n",
    "plt.ylabel(\"$y_{corr}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 40536 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-1, 1, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_9410_ty\"]),\n",
    "    y=ak.to_numpy(sel_array[\"yDiffOut\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"ty (T)\")\n",
    "plt.ylabel(\"$y_{corr}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 47128 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-0.5, 0.5, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_770_ty\"]),\n",
    "    y=ak.to_numpy(sel_array[\"yDiffOut\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"$t_y$\")\n",
    "plt.ylabel(\"$y_{corr}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/bend_y_ty_dist.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 40788 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-0.5, 0.5, 100)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"dSlope_yEndT\"]),\n",
    "    y=ak.to_numpy(sel_array[\"yDiffOut\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"$\\Delta t_y$\")\n",
    "plt.ylabel(\"$y_{corr}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/bend_y_deltaty_dist.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 41032 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-1.0, 1.0, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"dSlope_xEndT\"]),\n",
    "    y=ak.to_numpy(sel_array[\"yDiffOut\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"$\\Delta t_x$\")\n",
    "plt.ylabel(\"$y_{corr}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/bend_y_deltatx_dist.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "intercept= 0.0\n",
      "coef= {'dSlope_yEndT': 4089.1594362560113, 'dSlope_xEndT dSlope_yEndT': 25.0456971896117, 'ideal_state_770_ty dSlope_xEndT^2': 1049.7443418962382, 'ideal_state_770_ty dSlope_yEndT^2': 77388.96417801932}\n",
      "r2 score= 0.7591964011096296\n",
      "RMSE = 9.026349818744066\n"
     ]
    }
   ],
   "source": [
    "features = [\n",
    "    \"ideal_state_770_ty\",\n",
    "    # \"ideal_state_9410_ty\",\n",
    "    # \"ideal_state_770_y\",\n",
    "    # \"ideal_state_9410_y\",\n",
    "    \"dSlope_xEndT\",\n",
    "    \"dSlope_yEndT\",\n",
    "]\n",
    "target_feat = \"yDiffOut\"\n",
    "order = 3\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(sel_array[feat]) for feat in features])\n",
    "target = ak.to_numpy(sel_array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data,\n",
    "                                                    target,\n",
    "                                                    test_size=0.2,\n",
    "                                                    random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=order, 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",
    "\n",
    "# keep = [\n",
    "#     \"ideal_state_770_ty\",  #\n",
    "#     \"dSlope_xEndT\",  #\n",
    "#     \"dSlope_yEndT\",  # keep\n",
    "#     \"ideal_state_770_ty^2\",  #\n",
    "#     \"ideal_state_770_ty dSlope_xEndT\",  #\n",
    "#     \"ideal_state_770_ty dSlope_yEndT\",  #\n",
    "#     \"dSlope_xEndT^2\",  # do not keep\n",
    "#     \"dSlope_xEndT dSlope_yEndT\",  # keep\n",
    "#     \"dSlope_yEndT^2\",  #\n",
    "#     \"ideal_state_770_ty^3\",\n",
    "#     \"ideal_state_770_ty^2 dSlope_xEndT\",  #\n",
    "#     \"ideal_state_770_ty^2 dSlope_yEndT\",\n",
    "#     \"ideal_state_770_ty dSlope_xEndT^2\",  # keep\n",
    "#     \"ideal_state_770_ty dSlope_xEndT dSlope_yEndT\",  #\n",
    "#     \"ideal_state_770_ty dSlope_yEndT^2\",  # keep\n",
    "#     \"dSlope_xEndT^3\",  # do not keep\n",
    "#     \"dSlope_xEndT^2 dSlope_yEndT\",  #\n",
    "#     \"dSlope_xEndT dSlope_yEndT^2\",  #\n",
    "#     \"dSlope_yEndT^3\",\n",
    "#     \"ideal_state_770_ty^4\",\n",
    "#     \"ideal_state_770_ty^3 dSlope_xEndT\",\n",
    "#     \"ideal_state_770_ty^3 dSlope_yEndT\",\n",
    "#     \"ideal_state_770_ty^2 dSlope_xEndT^2\",  #\n",
    "#     \"ideal_state_770_ty^2 dSlope_xEndT dSlope_yEndT\",\n",
    "#     \"ideal_state_770_ty^2 dSlope_yEndT^2\",\n",
    "#     \"ideal_state_770_ty dSlope_xEndT^3\",  # do not keep\n",
    "#     \"ideal_state_770_ty dSlope_xEndT^2 dSlope_yEndT\",  #\n",
    "#     \"ideal_state_770_ty dSlope_xEndT dSlope_yEndT^2\",\n",
    "#     \"ideal_state_770_ty dSlope_yEndT^3\",\n",
    "#     \"dSlope_xEndT^4\",  #\n",
    "#     \"dSlope_xEndT^3 dSlope_yEndT\",  # keep\n",
    "#     \"dSlope_xEndT^2 dSlope_yEndT^2\",  #\n",
    "#     \"dSlope_xEndT dSlope_yEndT^3\",\n",
    "#     \"dSlope_yEndT^4\",\n",
    "#     \"ideal_state_770_ty^5\",\n",
    "#     \"ideal_state_770_ty^4 dSlope_xEndT\",\n",
    "#     \"ideal_state_770_ty^4 dSlope_yEndT\",\n",
    "#     \"ideal_state_770_ty^3 dSlope_xEndT^2\",\n",
    "#     \"ideal_state_770_ty^3 dSlope_xEndT dSlope_yEndT\",\n",
    "#     \"ideal_state_770_ty^3 dSlope_yEndT^2\",\n",
    "#     \"ideal_state_770_ty^2 dSlope_xEndT^3\",  #\n",
    "#     \"ideal_state_770_ty^2 dSlope_xEndT^2 dSlope_yEndT\",\n",
    "#     \"ideal_state_770_ty^2 dSlope_xEndT dSlope_yEndT^2\",\n",
    "#     \"ideal_state_770_ty^2 dSlope_yEndT^3\",\n",
    "#     \"ideal_state_770_ty dSlope_xEndT^4\",  # keep\n",
    "#     \"ideal_state_770_ty dSlope_xEndT^3 dSlope_yEndT\",  #\n",
    "#     \"ideal_state_770_ty dSlope_xEndT^2 dSlope_yEndT^2\",\n",
    "#     \"ideal_state_770_ty dSlope_xEndT dSlope_yEndT^3\",\n",
    "#     \"ideal_state_770_ty dSlope_yEndT^4\",\n",
    "#     \"dSlope_xEndT^5\",  #\n",
    "#     \"dSlope_xEndT^4 dSlope_yEndT\",  # keep\n",
    "#     \"dSlope_xEndT^3 dSlope_yEndT^2\",  #\n",
    "#     \"dSlope_xEndT^2 dSlope_yEndT^3\",\n",
    "#     \"dSlope_xEndT dSlope_yEndT^4\",\n",
    "#     \"dSlope_yEndT^5\",\n",
    "# ]\n",
    "\n",
    "keep = [\n",
    "    \"dSlope_yEndT\",  # keep\n",
    "    \"dSlope_xEndT dSlope_yEndT\",  # keep\n",
    "    \"ideal_state_770_ty dSlope_xEndT^2\",  # keep\n",
    "    \"ideal_state_770_ty dSlope_yEndT^2\",  # keep\n",
    "    # \"dSlope_xEndT^3 dSlope_yEndT\",  # do not keep\n",
    "    # \"ideal_state_770_ty dSlope_xEndT^4\",  # keep\n",
    "    # \"dSlope_xEndT^4 dSlope_yEndT\",  # keep\n",
    "    ###\n",
    "    # \"ideal_state_770_ty^5\",\n",
    "    # \"dSlope_xEndT^5\",\n",
    "    # \"dSlope_yEndT^5\",\n",
    "]\n",
    "\n",
    "# keep = [\n",
    "#     \"ideal_state_770_ty dSlope_yEndT^2\",\n",
    "#     \"ideal_state_770_ty dSlope_xEndT^2\",\n",
    "# ]\n",
    "remove = [i for i, f in enumerate(poly_features) if f not in keep]\n",
    "\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",
    "lin_reg = LinearRegression(fit_intercept=False)\n",
    "# lin_reg = Lasso(fit_intercept=False, alpha=0.000001)\n",
    "# lin_reg = Lasso(alpha=0.1)\n",
    "# lin_reg = LassoCV(max_iter=2000)\n",
    "# lin_reg = ElasticNet(alpha=0.1)\n",
    "# lin_reg = Ridge(alpha=0.1)\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))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 115,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "koeffs = dict(zip(poly_features, lin_reg.coef_))\n",
    "\n",
    "remove = []\n",
    "for itr in koeffs.items():\n",
    "    if abs(itr[1]) < 1.0:\n",
    "        remove.append(itr[0])\n",
    "remove"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['dSlope_yEndT',\n",
       " 'dSlope_xEndT dSlope_yEndT',\n",
       " 'ideal_state_770_ty dSlope_xEndT^2',\n",
       " 'ideal_state_770_ty dSlope_yEndT^2',\n",
       " 'ideal_state_770_ty dSlope_xEndT^4',\n",
       " 'dSlope_xEndT^4 dSlope_yEndT']"
      ]
     },
     "execution_count": 116,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ls_koeffs = []\n",
    "for itr in koeffs.items():\n",
    "    ls_koeffs.append(itr[0])\n",
    "ls_koeffs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 78860 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-40, 40, 41)\n",
    "ax = sns.regplot(\n",
    "    x=y_test,\n",
    "    y=y_test - y_pred_test,\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    "    label=\"bla\",\n",
    ")\n",
    "ax2 = ax.twinx()\n",
    "ax2.hist(\n",
    "    y_test,\n",
    "    bins=30,\n",
    "    range=[-40, 40],\n",
    "    color=\"#2A9D8F\",\n",
    "    alpha=0.5,\n",
    "    align=\"mid\",\n",
    "    density=True,\n",
    ")\n",
    "ax.set_xlabel(r\"$y_{\\mathrm{corr}}$ [mm]\")\n",
    "ax.set_ylabel(\n",
    "    r\"$\\left\\langle y_{\\mathrm{corr}}-y_{\\mathrm{corr}}^{\\mathrm{pred}}\\right\\rangle$ [mm]\"\n",
    ")\n",
    "ax2.set_ylabel(\"Number of Tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\", loc=0)\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/bend_y_regression_plot.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = [\n",
    "    \"ideal_state_770_ty\",\n",
    "    \"ideal_state_770_tx\",\n",
    "    \"dSlope_xEndT\",\n",
    "    \"dSlope_yEndT\",\n",
    "    \"dSlope_xEndT_abs\",\n",
    "    \"dSlope_yEndT_abs\",\n",
    "]\n",
    "target_feat = \"yDiffEndT\"\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,\n",
    "                                                    target,\n",
    "                                                    test_size=0.2,\n",
    "                                                    random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=6, 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",
    "    #'dSlope_xEndT',\n",
    "    \"dSlope_yEndT\",  # keep\n",
    "    #'dSlope_yEndT_abs',\n",
    "    #'ty dSlope_xEndT',\n",
    "    #'ty dSlope_yEndT',\n",
    "    \"ideal_state_770_ty dSlope_xEndT_abs\",  # keep\n",
    "    \"ideal_state_770_ty dSlope_yEndT_abs\",  # keep\n",
    "    \"ideal_state_770_ty dSlope_yEndT^2\",  # keep\n",
    "    \"ideal_state_770_ty dSlope_xEndT^2\",  # keep\n",
    "    #'tx dSlope_xEndT',\n",
    "    #'tx dSlope_xEndT_abs',\n",
    "    #'tx dSlope_yEndT',\n",
    "    \"ideal_state_770_ty ideal_state_770_tx dSlope_xEndT\",  # keep\n",
    "    \"ideal_state_770_tx^2 dSlope_yEndT\",  # keep\n",
    "    #'ty^2 dSlope_xEndT',\n",
    "    #'ty^2 dSlope_yEndT',\n",
    "    #'ty^2 dSlope_xEndT_abs',\n",
    "    #'ty^2 tx dSlope_xEndT',\n",
    "    #'ty tx^2 dSlope_yEndT',\n",
    "    \"ideal_state_770_ty ideal_state_770_tx^2 dSlope_xEndT_abs\",  # keep\n",
    "    \"ideal_state_770_ty^3 ideal_state_770_tx dSlope_xEndT\",  # keep\n",
    "    #'ty tx^3 dSlope_xEndT',\n",
    "    #'ty^3 dSlope_yEndT_abs',\n",
    "]\n",
    "do_not_keep = [\n",
    "    \"dSlope_xEndT\",\n",
    "    \"dSlope_yEndT_abs\",\n",
    "    \"ideal_state_770_ty dSlope_xEndT\",\n",
    "    \"ideal_state_770_tx dSlope_xEndT\",\n",
    "    \"ideal_state_770_tx dSlope_xEndT_abs\",\n",
    "    \"ideal_state_770_tx dSlope_yEndT\",\n",
    "    \"ideal_state_770_ty^2 dSlope_xEndT\",\n",
    "    \"ideal_state_770_ty^3 dSlope_yEndT_abs\",\n",
    "    \"ideal_state_770_ty ideal_state_770_tx dSlope_yEndT\",\n",
    "    \"ideal_state_770_ty ideal_state_770_tx^3 dSlope_xEndT\",\n",
    "    \"ideal_state_770_ty ideal_state_770_tx^2 dSlope_yEndT\",\n",
    "]\n",
    "reduce = True\n",
    "if reduce:\n",
    "    remove = [\n",
    "        i for i, f in enumerate(poly_features) if (keep and f not in keep)\n",
    "    ]\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",
    "if not reduce:\n",
    "    remove = [\n",
    "        i for i, f in enumerate(poly_features)\n",
    "        if (\"dSlope_\" not in f) or (\"EndT^\" in f) or (\"abs^\" in f) or (\n",
    "            \"EndT dSlope\" in f) or (\"abs dSlope\" in f)\n",
    "    ]\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",
    "    lin_reg = Lasso(fit_intercept=False, alpha=0.000001)\n",
    "else:\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(\n",
    "    \"coef=\",\n",
    "    {\n",
    "        k: v\n",
    "        for k, v in zip(poly_features, lin_reg.coef_)\n",
    "        if abs(v) > 1.0 and k not in keep and k not in do_not_keep\n",
    "    },\n",
    ")\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(\n",
    "    \"straight RMSE =\",\n",
    "    mean_squared_error(\n",
    "        array[\"ideal_state_9410_y\"],\n",
    "        array[\"ideal_state_770_y\"] + array[\"ideal_state_770_ty\"] *\n",
    "        (array[\"ideal_state_9410_z\"] - array[\"ideal_state_770_z\"]),\n",
    "        squared=False,\n",
    "    ),\n",
    ")\n",
    "print(format_array(\"y_xEndT_diff\", lin_reg.coef_))"
   ]
  }
 ],
 "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.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "a2eff8b4da8b8eebf5ee2e5f811f31a557e0a202b4d2f04f849b065340a6eda6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
