{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d563033c",
   "metadata": {},
   "source": [
    "# Import tools/models be used"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "id": "fc9e6e26",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from math import pi\n",
    "from sklearn.linear_model import LinearRegression     # import linear model to be used\n",
    "from sklearn.linear_model import SGDRegressor     # import linear model to be used\n",
    "import joblib                                         # to store model\n",
    "import os\n",
    "import errno\n",
    "from sklearn.model_selection import train_test_split  # this is very important"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "id": "79a93533",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    os.mkdir('Plots')\n",
    "# if folder exists do nothing\n",
    "except OSError as e:\n",
    "    if e.errno != errno.EEXIST:\n",
    "        raise\n",
    "        \n",
    "try:\n",
    "    os.mkdir('storage')\n",
    "# if folder exists do nothing\n",
    "except OSError as e:\n",
    "    if e.errno != errno.EEXIST:\n",
    "        raise        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "313e157e",
   "metadata": {},
   "source": [
    "# Read data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "id": "bdc08ed0",
   "metadata": {},
   "outputs": [
    {
     "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>$M_{inv}$</th>\n",
       "      <th>$rapidity$</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>512.889687</td>\n",
       "      <td>-1.698364</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>349.604335</td>\n",
       "      <td>0.649795</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>102.015218</td>\n",
       "      <td>0.935912</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>475.693746</td>\n",
       "      <td>1.691695</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>422.624040</td>\n",
       "      <td>1.233733</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8604</th>\n",
       "      <td>310.392511</td>\n",
       "      <td>1.763102</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8605</th>\n",
       "      <td>322.772764</td>\n",
       "      <td>1.343101</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8606</th>\n",
       "      <td>373.867518</td>\n",
       "      <td>-0.255696</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8607</th>\n",
       "      <td>275.877980</td>\n",
       "      <td>0.169265</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8608</th>\n",
       "      <td>440.485176</td>\n",
       "      <td>0.265116</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>8609 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "       $M_{inv}$  $rapidity$\n",
       "0     512.889687   -1.698364\n",
       "1     349.604335    0.649795\n",
       "2     102.015218    0.935912\n",
       "3     475.693746    1.691695\n",
       "4     422.624040    1.233733\n",
       "...          ...         ...\n",
       "8604  310.392511    1.763102\n",
       "8605  322.772764    1.343101\n",
       "8606  373.867518   -0.255696\n",
       "8607  275.877980    0.169265\n",
       "8608  440.485176    0.265116\n",
       "\n",
       "[8609 rows x 2 columns]"
      ]
     },
     "execution_count": 145,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "directory = '../data/'\n",
    "filein    = 'checkNLO.txt'\n",
    "file      = directory+filein\n",
    "data = np.genfromtxt(file)\n",
    "\n",
    "m        = data[:,0:1]\n",
    "rap      = data[:,1:2]\n",
    "x2       = data[:,2:3]\n",
    "\n",
    "m = np.concatenate((m),axis=0) \n",
    "rap = np.concatenate((rap),axis=0) \n",
    "x2 = np.concatenate((x2),axis=0) \n",
    "\n",
    "Data_ave = np.stack((m,rap),axis=-1)\n",
    "df = pd.DataFrame(Data_ave)\n",
    "df.columns = ['$M_{inv}$','$rapidity$']\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0248b96b",
   "metadata": {},
   "source": [
    "# Split data into training and testing sub-sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "id": "de4cf784",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for reproducibility\n",
    "#X_train, X_test, y_train, y_test = train_test_split(Data_ave, x2, test_size=0.4, train_size=0.6 , random_state=0)\n",
    "X_train, X_test, y_train, y_test = train_test_split(Data_ave, x2, test_size=0.4, train_size=0.6 )\n",
    "\n",
    "# rename data for convenience\n",
    "x2_train   = y_train\n",
    "m_train    = X_train[:,0:1]\n",
    "rap_train  = X_train[:,1:2]\n",
    "\n",
    "x2_test   = y_test\n",
    "m_test    = X_test[:,0:1]\n",
    "rap_test  = X_test[:,1:2]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2624e86c",
   "metadata": {},
   "source": [
    "# Now let's do some training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "id": "b8ef444a",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''Training and prediction with Linear Model'''\n",
    "#-----------------------------------------------------------------------------#\n",
    "# Input:                                                                      #\n",
    "# X_train  := training data set                                               #\n",
    "# y_train  := targets for training                                            #\n",
    "# X_trest  := testing data set                                                #\n",
    "# y_test   := targets for testing                                             #\n",
    "# filename := name for training storage                                       #\n",
    "# model    := name of the mode we want to use                                 #\n",
    "# ----------------------------------------------------------------------------#\n",
    "# Output:                                                                     #  \n",
    "# loaded_model := trained model                                               #\n",
    "# R2 := quality of the training on the testing set                            #\n",
    "# y_fit := predicted targets                                                  #\n",
    "# y_est := estimated targets                                                  #\n",
    "# ----------------------------------------------------------------------------#\n",
    "# Output labelling :                                                          #\n",
    "# Training(X,y,filename)[0] -> R2                                             #\n",
    "# Training(X,y,filename)[1] -> y_new                                          #\n",
    "# ----------------------------------------------------------------------------#\n",
    "\n",
    "def training (X_train, y_train, X_test, y_test,filename,model):\n",
    "\n",
    "    # train \n",
    "    model.fit(X_train, y_train)\n",
    "\n",
    "    # save trained model\n",
    "    joblib.dump(model, filename)\n",
    "\n",
    "    # load model \n",
    "    loaded_model = joblib.load(filename)\n",
    "\n",
    "    # R^2\n",
    "    R2   = loaded_model.score(X_test, y_test)\n",
    "    \n",
    "    # predicted targets \n",
    "    y_fit = loaded_model.predict(X_test)\n",
    "    # estimated targets\n",
    "    y_est = loaded_model.predict(X_train)\n",
    "\n",
    "    \n",
    "    return R2, y_fit,y_est, model.coef_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "id": "be16d388",
   "metadata": {},
   "outputs": [],
   "source": [
    "# name of the model for storage\n",
    "filename = 'storage/x2_NLO_LinearRegression_model_not_thinking.sav'\n",
    "\n",
    "# train \n",
    "Training_machine_learning = training(X_train,y_train,X_test,y_test,filename,LinearRegression())\n",
    "\n",
    "# predict\n",
    "R2    = Training_machine_learning[0] \n",
    "y_fit = Training_machine_learning[1]\n",
    "y_est = Training_machine_learning[2]\n",
    "coefs = Training_machine_learning[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 191,
   "id": "93efe92e",
   "metadata": {},
   "outputs": [],
   "source": [
    "coefs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "id": "9c5a08b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "R2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56253913",
   "metadata": {},
   "source": [
    "# Now, let's use Stochastic Gradient Descent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "id": "51451285",
   "metadata": {},
   "outputs": [],
   "source": [
    "# name of the model for storage\n",
    "filename = 'storage/x2_NLO_SGD_model_not_thinking.sav'\n",
    "\n",
    "# train \n",
    "Training_machine_learning = training(X_train,y_train,X_test,y_test,filename,\n",
    "                                     SGDRegressor(max_iter=1e9,tol=1e-3,penalty=None,eta0=1e-6))\n",
    "\n",
    "# predict\n",
    "R2    = Training_machine_learning[0] \n",
    "y_fit = Training_machine_learning[1]\n",
    "y_est = Training_machine_learning[2]\n",
    "coefs = Training_machine_learning[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 188,
   "id": "fc9e5f11",
   "metadata": {},
   "outputs": [],
   "source": [
    "R2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 189,
   "id": "c51df2a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "coefs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e69974b",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}