notebooks/Logistic_regression.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4c5bbe66-3340-4f47-bf62-3a486b573b2a",
   "metadata": {},
   "source": [
    "# Logistic Regression "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "71b53af7-9669-47a5-95d4-b574afcf4627",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tenseal as ts\n",
    "import pandas as pd\n",
    "import random\n",
    "import torch\n",
    "from torch import nn\n",
    "from time import time\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.metrics import classification_report, confusion_matrix, accuracy_score"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3248e6f0-d13a-4e3b-8405-786642ba0055",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "\n",
    "Import the dataset from Kaggle, you can download the data set [here](https://www.kaggle.com/datasets/naveengowda16/logistic-regression-heart-disease-prediction?resource=download). This data set contains data by WHO regarding heart disease. \n",
    "\n",
    "World Health Organization has estimated 12 million deaths occur worldwide, every year due to Heart diseases. Half the deaths in the United States and other developed countries are due to cardio vascular diseases. The early prognosis of cardiovascular diseases can aid in making decisions on lifestyle changes in high risk patients and in turn reduce the complications. This research intends to pinpoint the most relevant/risk factors of heart disease as well as predict the overall risk using logistic regression.\n",
    "\n",
    "This dataset includes patients' information along with a 10-year risk of future coronary heart disease (CHD) as a label. The goal is to build a model that can predict this 10-year CHD risk based on patients' information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "215ca018-2595-4265-af04-c1aa830946d1",
   "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>male</th>\n",
       "      <th>age</th>\n",
       "      <th>education</th>\n",
       "      <th>currentSmoker</th>\n",
       "      <th>cigsPerDay</th>\n",
       "      <th>BPMeds</th>\n",
       "      <th>prevalentStroke</th>\n",
       "      <th>prevalentHyp</th>\n",
       "      <th>diabetes</th>\n",
       "      <th>totChol</th>\n",
       "      <th>sysBP</th>\n",
       "      <th>diaBP</th>\n",
       "      <th>BMI</th>\n",
       "      <th>heartRate</th>\n",
       "      <th>glucose</th>\n",
       "      <th>TenYearCHD</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>39</td>\n",
       "      <td>4.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>195.0</td>\n",
       "      <td>106.0</td>\n",
       "      <td>70.0</td>\n",
       "      <td>26.97</td>\n",
       "      <td>80.0</td>\n",
       "      <td>77.0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>46</td>\n",
       "      <td>2.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>250.0</td>\n",
       "      <td>121.0</td>\n",
       "      <td>81.0</td>\n",
       "      <td>28.73</td>\n",
       "      <td>95.0</td>\n",
       "      <td>76.0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>48</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>20.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>245.0</td>\n",
       "      <td>127.5</td>\n",
       "      <td>80.0</td>\n",
       "      <td>25.34</td>\n",
       "      <td>75.0</td>\n",
       "      <td>70.0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>61</td>\n",
       "      <td>3.0</td>\n",
       "      <td>1</td>\n",
       "      <td>30.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>225.0</td>\n",
       "      <td>150.0</td>\n",
       "      <td>95.0</td>\n",
       "      <td>28.58</td>\n",
       "      <td>65.0</td>\n",
       "      <td>103.0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>46</td>\n",
       "      <td>3.0</td>\n",
       "      <td>1</td>\n",
       "      <td>23.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>285.0</td>\n",
       "      <td>130.0</td>\n",
       "      <td>84.0</td>\n",
       "      <td>23.10</td>\n",
       "      <td>85.0</td>\n",
       "      <td>85.0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   male  age  education  currentSmoker  cigsPerDay  BPMeds  prevalentStroke  \\\n",
       "0     1   39        4.0              0         0.0     0.0                0   \n",
       "1     0   46        2.0              0         0.0     0.0                0   \n",
       "2     1   48        1.0              1        20.0     0.0                0   \n",
       "3     0   61        3.0              1        30.0     0.0                0   \n",
       "4     0   46        3.0              1        23.0     0.0                0   \n",
       "\n",
       "   prevalentHyp  diabetes  totChol  sysBP  diaBP    BMI  heartRate  glucose  \\\n",
       "0             0         0    195.0  106.0   70.0  26.97       80.0     77.0   \n",
       "1             0         0    250.0  121.0   81.0  28.73       95.0     76.0   \n",
       "2             0         0    245.0  127.5   80.0  25.34       75.0     70.0   \n",
       "3             1         0    225.0  150.0   95.0  28.58       65.0    103.0   \n",
       "4             0         0    285.0  130.0   84.0  23.10       85.0     85.0   \n",
       "\n",
       "   TenYearCHD  \n",
       "0           0  \n",
       "1           0  \n",
       "2           0  \n",
       "3           1  \n",
       "4           0  "
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv(\"../data/raw/framingham_heart_disease.csv\")\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b29b65b9-bc66-40d3-8c5b-9660122db3f9",
   "metadata": {},
   "source": [
    "## EDA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "88fc6daa-2cb1-4661-8a56-4d38f7cc912e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4238, 16)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5c3dd232-9119-40be-bb74-eda109918ec1",
   "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>male</th>\n",
       "      <th>age</th>\n",
       "      <th>education</th>\n",
       "      <th>currentSmoker</th>\n",
       "      <th>cigsPerDay</th>\n",
       "      <th>BPMeds</th>\n",
       "      <th>prevalentStroke</th>\n",
       "      <th>prevalentHyp</th>\n",
       "      <th>diabetes</th>\n",
       "      <th>totChol</th>\n",
       "      <th>sysBP</th>\n",
       "      <th>diaBP</th>\n",
       "      <th>BMI</th>\n",
       "      <th>heartRate</th>\n",
       "      <th>glucose</th>\n",
       "      <th>TenYearCHD</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>count</th>\n",
       "      <td>4238.000000</td>\n",
       "      <td>4238.000000</td>\n",
       "      <td>4133.000000</td>\n",
       "      <td>4238.000000</td>\n",
       "      <td>4209.000000</td>\n",
       "      <td>4185.000000</td>\n",
       "      <td>4238.000000</td>\n",
       "      <td>4238.000000</td>\n",
       "      <td>4238.000000</td>\n",
       "      <td>4188.000000</td>\n",
       "      <td>4238.000000</td>\n",
       "      <td>4238.000000</td>\n",
       "      <td>4219.000000</td>\n",
       "      <td>4237.000000</td>\n",
       "      <td>3850.000000</td>\n",
       "      <td>4238.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.429212</td>\n",
       "      <td>49.584946</td>\n",
       "      <td>1.978950</td>\n",
       "      <td>0.494101</td>\n",
       "      <td>9.003089</td>\n",
       "      <td>0.029630</td>\n",
       "      <td>0.005899</td>\n",
       "      <td>0.310524</td>\n",
       "      <td>0.025720</td>\n",
       "      <td>236.721585</td>\n",
       "      <td>132.352407</td>\n",
       "      <td>82.893464</td>\n",
       "      <td>25.802008</td>\n",
       "      <td>75.878924</td>\n",
       "      <td>81.966753</td>\n",
       "      <td>0.151958</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.495022</td>\n",
       "      <td>8.572160</td>\n",
       "      <td>1.019791</td>\n",
       "      <td>0.500024</td>\n",
       "      <td>11.920094</td>\n",
       "      <td>0.169584</td>\n",
       "      <td>0.076587</td>\n",
       "      <td>0.462763</td>\n",
       "      <td>0.158316</td>\n",
       "      <td>44.590334</td>\n",
       "      <td>22.038097</td>\n",
       "      <td>11.910850</td>\n",
       "      <td>4.080111</td>\n",
       "      <td>12.026596</td>\n",
       "      <td>23.959998</td>\n",
       "      <td>0.359023</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>min</th>\n",
       "      <td>0.000000</td>\n",
       "      <td>32.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>107.000000</td>\n",
       "      <td>83.500000</td>\n",
       "      <td>48.000000</td>\n",
       "      <td>15.540000</td>\n",
       "      <td>44.000000</td>\n",
       "      <td>40.000000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25%</th>\n",
       "      <td>0.000000</td>\n",
       "      <td>42.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>206.000000</td>\n",
       "      <td>117.000000</td>\n",
       "      <td>75.000000</td>\n",
       "      <td>23.070000</td>\n",
       "      <td>68.000000</td>\n",
       "      <td>71.000000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50%</th>\n",
       "      <td>0.000000</td>\n",
       "      <td>49.000000</td>\n",
       "      <td>2.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>234.000000</td>\n",
       "      <td>128.000000</td>\n",
       "      <td>82.000000</td>\n",
       "      <td>25.400000</td>\n",
       "      <td>75.000000</td>\n",
       "      <td>78.000000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>75%</th>\n",
       "      <td>1.000000</td>\n",
       "      <td>56.000000</td>\n",
       "      <td>3.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>20.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>263.000000</td>\n",
       "      <td>144.000000</td>\n",
       "      <td>89.875000</td>\n",
       "      <td>28.040000</td>\n",
       "      <td>83.000000</td>\n",
       "      <td>87.000000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>max</th>\n",
       "      <td>1.000000</td>\n",
       "      <td>70.000000</td>\n",
       "      <td>4.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>70.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>696.000000</td>\n",
       "      <td>295.000000</td>\n",
       "      <td>142.500000</td>\n",
       "      <td>56.800000</td>\n",
       "      <td>143.000000</td>\n",
       "      <td>394.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              male          age    education  currentSmoker   cigsPerDay  \\\n",
       "count  4238.000000  4238.000000  4133.000000    4238.000000  4209.000000   \n",
       "mean      0.429212    49.584946     1.978950       0.494101     9.003089   \n",
       "std       0.495022     8.572160     1.019791       0.500024    11.920094   \n",
       "min       0.000000    32.000000     1.000000       0.000000     0.000000   \n",
       "25%       0.000000    42.000000     1.000000       0.000000     0.000000   \n",
       "50%       0.000000    49.000000     2.000000       0.000000     0.000000   \n",
       "75%       1.000000    56.000000     3.000000       1.000000    20.000000   \n",
       "max       1.000000    70.000000     4.000000       1.000000    70.000000   \n",
       "\n",
       "            BPMeds  prevalentStroke  prevalentHyp     diabetes      totChol  \\\n",
       "count  4185.000000      4238.000000   4238.000000  4238.000000  4188.000000   \n",
       "mean      0.029630         0.005899      0.310524     0.025720   236.721585   \n",
       "std       0.169584         0.076587      0.462763     0.158316    44.590334   \n",
       "min       0.000000         0.000000      0.000000     0.000000   107.000000   \n",
       "25%       0.000000         0.000000      0.000000     0.000000   206.000000   \n",
       "50%       0.000000         0.000000      0.000000     0.000000   234.000000   \n",
       "75%       0.000000         0.000000      1.000000     0.000000   263.000000   \n",
       "max       1.000000         1.000000      1.000000     1.000000   696.000000   \n",
       "\n",
       "             sysBP        diaBP          BMI    heartRate      glucose  \\\n",
       "count  4238.000000  4238.000000  4219.000000  4237.000000  3850.000000   \n",
       "mean    132.352407    82.893464    25.802008    75.878924    81.966753   \n",
       "std      22.038097    11.910850     4.080111    12.026596    23.959998   \n",
       "min      83.500000    48.000000    15.540000    44.000000    40.000000   \n",
       "25%     117.000000    75.000000    23.070000    68.000000    71.000000   \n",
       "50%     128.000000    82.000000    25.400000    75.000000    78.000000   \n",
       "75%     144.000000    89.875000    28.040000    83.000000    87.000000   \n",
       "max     295.000000   142.500000    56.800000   143.000000   394.000000   \n",
       "\n",
       "        TenYearCHD  \n",
       "count  4238.000000  \n",
       "mean      0.151958  \n",
       "std       0.359023  \n",
       "min       0.000000  \n",
       "25%       0.000000  \n",
       "50%       0.000000  \n",
       "75%       0.000000  \n",
       "max       1.000000  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "77c6c99a-fc03-4cfc-b486-42d67eabfe2c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "male                 0\n",
       "age                  0\n",
       "education          105\n",
       "currentSmoker        0\n",
       "cigsPerDay          29\n",
       "BPMeds              53\n",
       "prevalentStroke      0\n",
       "prevalentHyp         0\n",
       "diabetes             0\n",
       "totChol             50\n",
       "sysBP                0\n",
       "diaBP                0\n",
       "BMI                 19\n",
       "heartRate            1\n",
       "glucose            388\n",
       "TenYearCHD           0\n",
       "dtype: int64"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check for any missing values\n",
    "df.isna().sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fda230fc-29b3-4a1a-8cb0-c5e76a31d69d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def clean_dataset(df):\n",
    "\n",
    "    # We could either populate the missing values or just\n",
    "    # remove the entire column if not needed any more.\n",
    "\n",
    "    df[\"education\"] = df[\"education\"].fillna(df[\"education\"].mean())\n",
    "    df[\"cigsPerDay\"] = df[\"cigsPerDay\"].fillna(df[\"cigsPerDay\"].mean())\n",
    "    df[\"BPMeds\"] = df[\"BPMeds\"].fillna(df[\"BPMeds\"].mean())\n",
    "    df[\"totChol\"] = df[\"totChol\"].fillna(df[\"totChol\"].mean())\n",
    "    df[\"BMI\"] = df[\"BMI\"].fillna(df[\"BMI\"].mean())\n",
    "    df[\"heartRate\"] = df[\"heartRate\"].fillna(df[\"heartRate\"].mean())\n",
    "    df[\"glucose\"] = df[\"glucose\"].fillna(df[\"glucose\"].mean())\n",
    "\n",
    "    # remove rows with missing values\n",
    "    df = df.dropna()\n",
    "\n",
    "    new_df = df.groupby(\"TenYearCHD\")\n",
    "    df = new_df.apply(lambda x: x.sample(new_df.size().min(), random_state=73).reset_index(drop=True))\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e89f6083-cde8-4dad-883c-541cf045533e",
   "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></th>\n",
       "      <th>male</th>\n",
       "      <th>age</th>\n",
       "      <th>education</th>\n",
       "      <th>currentSmoker</th>\n",
       "      <th>cigsPerDay</th>\n",
       "      <th>BPMeds</th>\n",
       "      <th>prevalentStroke</th>\n",
       "      <th>prevalentHyp</th>\n",
       "      <th>diabetes</th>\n",
       "      <th>totChol</th>\n",
       "      <th>sysBP</th>\n",
       "      <th>diaBP</th>\n",
       "      <th>BMI</th>\n",
       "      <th>heartRate</th>\n",
       "      <th>glucose</th>\n",
       "      <th>TenYearCHD</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>TenYearCHD</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"5\" valign=\"top\">0</th>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>44</td>\n",
       "      <td>3.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>243.0</td>\n",
       "      <td>146.0</td>\n",
       "      <td>91.0</td>\n",
       "      <td>26.72</td>\n",
       "      <td>80.0</td>\n",
       "      <td>104.000000</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>60</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>236.0</td>\n",
       "      <td>126.0</td>\n",
       "      <td>84.0</td>\n",
       "      <td>20.34</td>\n",
       "      <td>71.0</td>\n",
       "      <td>76.000000</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>45</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>30.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>240.0</td>\n",
       "      <td>141.0</td>\n",
       "      <td>89.0</td>\n",
       "      <td>25.01</td>\n",
       "      <td>95.0</td>\n",
       "      <td>76.000000</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>47</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>30.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>190.0</td>\n",
       "      <td>147.5</td>\n",
       "      <td>92.5</td>\n",
       "      <td>31.31</td>\n",
       "      <td>77.0</td>\n",
       "      <td>82.000000</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>43</td>\n",
       "      <td>4.0</td>\n",
       "      <td>1</td>\n",
       "      <td>9.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>207.0</td>\n",
       "      <td>95.5</td>\n",
       "      <td>70.0</td>\n",
       "      <td>19.78</td>\n",
       "      <td>93.0</td>\n",
       "      <td>79.000000</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"5\" valign=\"top\">1</th>\n",
       "      <th>639</th>\n",
       "      <td>1</td>\n",
       "      <td>62</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>43.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>217.0</td>\n",
       "      <td>107.5</td>\n",
       "      <td>75.0</td>\n",
       "      <td>26.21</td>\n",
       "      <td>80.0</td>\n",
       "      <td>66.000000</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>640</th>\n",
       "      <td>1</td>\n",
       "      <td>37</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>30.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>179.0</td>\n",
       "      <td>125.0</td>\n",
       "      <td>82.0</td>\n",
       "      <td>19.53</td>\n",
       "      <td>60.0</td>\n",
       "      <td>70.000000</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>641</th>\n",
       "      <td>1</td>\n",
       "      <td>59</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>3.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>230.0</td>\n",
       "      <td>182.0</td>\n",
       "      <td>102.0</td>\n",
       "      <td>25.91</td>\n",
       "      <td>66.0</td>\n",
       "      <td>147.000000</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>642</th>\n",
       "      <td>0</td>\n",
       "      <td>58</td>\n",
       "      <td>3.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>241.0</td>\n",
       "      <td>143.5</td>\n",
       "      <td>85.5</td>\n",
       "      <td>23.96</td>\n",
       "      <td>96.0</td>\n",
       "      <td>81.966753</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>643</th>\n",
       "      <td>1</td>\n",
       "      <td>51</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>20.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>239.0</td>\n",
       "      <td>168.0</td>\n",
       "      <td>102.0</td>\n",
       "      <td>30.38</td>\n",
       "      <td>82.0</td>\n",
       "      <td>68.000000</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1288 rows × 16 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                male  age  education  currentSmoker  cigsPerDay  BPMeds  \\\n",
       "TenYearCHD                                                                \n",
       "0          0       1   44        3.0              0         0.0     0.0   \n",
       "           1       0   60        1.0              0         0.0     0.0   \n",
       "           2       1   45        1.0              1        30.0     0.0   \n",
       "           3       1   47        1.0              1        30.0     0.0   \n",
       "           4       0   43        4.0              1         9.0     0.0   \n",
       "...              ...  ...        ...            ...         ...     ...   \n",
       "1          639     1   62        1.0              1        43.0     0.0   \n",
       "           640     1   37        1.0              1        30.0     0.0   \n",
       "           641     1   59        1.0              1         3.0     0.0   \n",
       "           642     0   58        3.0              0         0.0     0.0   \n",
       "           643     1   51        1.0              1        20.0     0.0   \n",
       "\n",
       "                prevalentStroke  prevalentHyp  diabetes  totChol  sysBP  \\\n",
       "TenYearCHD                                                                \n",
       "0          0                  0             0         0    243.0  146.0   \n",
       "           1                  0             0         0    236.0  126.0   \n",
       "           2                  0             0         0    240.0  141.0   \n",
       "           3                  0             1         0    190.0  147.5   \n",
       "           4                  0             0         0    207.0   95.5   \n",
       "...                         ...           ...       ...      ...    ...   \n",
       "1          639                0             0         0    217.0  107.5   \n",
       "           640                0             0         0    179.0  125.0   \n",
       "           641                0             1         1    230.0  182.0   \n",
       "           642                0             1         0    241.0  143.5   \n",
       "           643                0             1         0    239.0  168.0   \n",
       "\n",
       "                diaBP    BMI  heartRate     glucose  TenYearCHD  \n",
       "TenYearCHD                                                       \n",
       "0          0     91.0  26.72       80.0  104.000000           0  \n",
       "           1     84.0  20.34       71.0   76.000000           0  \n",
       "           2     89.0  25.01       95.0   76.000000           0  \n",
       "           3     92.5  31.31       77.0   82.000000           0  \n",
       "           4     70.0  19.78       93.0   79.000000           0  \n",
       "...               ...    ...        ...         ...         ...  \n",
       "1          639   75.0  26.21       80.0   66.000000           1  \n",
       "           640   82.0  19.53       60.0   70.000000           1  \n",
       "           641  102.0  25.91       66.0  147.000000           1  \n",
       "           642   85.5  23.96       96.0   81.966753           1  \n",
       "           643  102.0  30.38       82.0   68.000000           1  \n",
       "\n",
       "[1288 rows x 16 columns]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "clean_dataset(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fea83fee-8d71-4353-a733-50d1902dd2d6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "male               0\n",
       "age                0\n",
       "education          0\n",
       "currentSmoker      0\n",
       "cigsPerDay         0\n",
       "BPMeds             0\n",
       "prevalentStroke    0\n",
       "prevalentHyp       0\n",
       "diabetes           0\n",
       "totChol            0\n",
       "sysBP              0\n",
       "diaBP              0\n",
       "BMI                0\n",
       "heartRate          0\n",
       "glucose            0\n",
       "TenYearCHD         0\n",
       "dtype: int64"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check if any null or na values exist\n",
    "df.isna().sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "936ad77f-11b2-4baf-9d69-1a40cb6b47d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fetch X and y\n",
    "y = df[\"TenYearCHD\"].values\n",
    "X = df.drop(\"TenYearCHD\", axis = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "71eecb49-d0bf-4ada-8612-c9605cfd1024",
   "metadata": {},
   "outputs": [],
   "source": [
    " X_train, X_test, y_train, y_test = train_test_split(X,y,random_state = 0,test_size = 0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f3a3602e-e023-4ba7-8276-5f80d1209187",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X_train has shape: (3390, 15)\n",
      "y_train has shape: (3390,)\n",
      "X_test has shape: (848, 15)\n",
      "y_test has shape: (848,)\n"
     ]
    }
   ],
   "source": [
    "print(f\"X_train has shape: {X_train.shape}\")\n",
    "print(f\"y_train has shape: {y_train.shape}\")\n",
    "print(f\"X_test has shape: {X_test.shape}\")\n",
    "print(f\"y_test has shape: {y_test.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d425bbb-8195-46e0-8dae-ac47f10a2c31",
   "metadata": {},
   "source": [
    "## Train a Logistic Regression Model using plaintext data\n",
    "\n",
    "We will start by training a logistic regression model (without any encryption), which can be viewed as a single layer neural network with a single node. We will be using this model as a means of comparison against encrypted training and evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f3919861-7295-44d3-9930-b8e32f0d83ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/app-root/lib64/python3.8/site-packages/sklearn/linear_model/_logistic.py:763: ConvergenceWarning: lbfgs failed to converge (status=1):\n",
      "STOP: TOTAL NO. of ITERATIONS REACHED LIMIT.\n",
      "\n",
      "Increase the number of iterations (max_iter) or scale the data as shown in:\n",
      "    https://scikit-learn.org/stable/modules/preprocessing.html\n",
      "Please also refer to the documentation for alternative solver options:\n",
      "    https://scikit-learn.org/stable/modules/linear_model.html#logistic-regression\n",
      "  n_iter_i = _check_optimize_result(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "LogisticRegression(random_state=0)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "classifier = LogisticRegression(random_state=0)\n",
    "classifier.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d6c8bfb3-db85-4e4b-9c7d-5d2d9f0baa2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = classifier.predict(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "5851bf59-a1d3-425d-a306-c9c381184838",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy :  0.8360849056603774\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.84      0.99      0.91       710\n",
      "           1       0.43      0.02      0.04       138\n",
      "\n",
      "    accuracy                           0.84       848\n",
      "   macro avg       0.63      0.51      0.48       848\n",
      "weighted avg       0.77      0.84      0.77       848\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"Accuracy : \", accuracy_score(y_test, y_pred))\n",
    "print(classification_report(y_test, y_pred))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51a5030d-11c7-4fb1-ae27-933994ee1fea",
   "metadata": {},
   "source": [
    "## Train a Logistic Regression Model \n",
    "We wil now train a Logistic Regression with a single neural network with a single node. \n",
    "\n",
    "[Reference](https://towardsdatascience.com/logistic-regression-with-pytorch-3c8bbea594be)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a51574a-2f81-4a5c-a538-925e00972a52",
   "metadata": {},
   "source": [
    "#### Preprocess data for torch model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "8afad501-91c7-4f03-b332-d81956bb0931",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X_train has shape: torch.Size([3391, 16])\n",
      "y_train has shape: torch.Size([3391, 1])\n",
      "X_test has shape: torch.Size([847, 16])\n",
      "y_test has shape: torch.Size([847, 1])\n"
     ]
    }
   ],
   "source": [
    "torch.random.manual_seed(73)\n",
    "random.seed(73)\n",
    "\n",
    "\n",
    "# Split the dataset into train(80%) and test(20%)\n",
    "def split_train_and_test(X, y, ratio=0.2):\n",
    "    idxs = [i for i in range(len(X))]\n",
    "    random.shuffle(idxs)\n",
    "    # delimiter between test and train data\n",
    "    delim = int(len(X) * ratio)\n",
    "    test_idxs, train_idxs = idxs[:delim], idxs[delim:]\n",
    "    return X[train_idxs], y[train_idxs], X[test_idxs], y[test_idxs]\n",
    "\n",
    "\n",
    "def standardize_data():\n",
    "    global df\n",
    "    clean_dataset(df)\n",
    "    df = (df - df.mean()) / df.std()\n",
    "    X = torch.tensor(df.values).float()\n",
    "    y = torch.tensor(df[\"TenYearCHD\"].values).float().unsqueeze(1)\n",
    "    df = df.drop(\"TenYearCHD\", 'columns')\n",
    "    return split_train_and_test(X, y)\n",
    "\n",
    "\n",
    "def randomize_data(m=1024, n=2):\n",
    "    # data separable by the line `y = x`\n",
    "    X_train = torch.randn(m, n)\n",
    "    X_test = torch.randn(m // 2, n)\n",
    "    y_train = (X_train[:, 0] >= X_train[:, 1]).float().unsqueeze(0).t()\n",
    "    y_test = (x_test[:, 0] >= X_test[:, 1]).float().unsqueeze(0).t()\n",
    "    return X_train, y_train, X_test, y_test\n",
    "\n",
    "X_train, y_train, X_test, y_test = standardize_data()\n",
    "\n",
    "\n",
    "print(f\"X_train has shape: {X_train.shape}\")\n",
    "print(f\"y_train has shape: {y_train.shape}\")\n",
    "print(f\"X_test has shape: {X_test.shape}\")\n",
    "print(f\"y_test has shape: {y_test.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "48c872a5-00df-48d9-aeb3-ad43a68ff31a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LogisticRegression(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, n_features):\n",
    "        super(LogisticRegression, self).__init__()\n",
    "        self.lr = torch.nn.Linear(n_features, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        output = torch.sigmoid(self.lr(x))\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "67a6b9e2-bab2-496b-9b82-b96277ca8963",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_features = X_train.shape[1]\n",
    "model = LogisticRegression(n_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "572561a5-1a70-40e2-8802-a1c0c05b53e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Gradient descent with learning rate of 1\n",
    "optim = torch.optim.SGD(model.parameters(), lr=1)\n",
    "bceloss = torch.nn.BCELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "51e5c282-4b35-4aae-88f9-541664a77316",
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f9698592-de9b-4574-8eea-74e28f45bc43",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, optim, bceloss, X, y, epochs):\n",
    "    for ep in range(1, epochs+1):\n",
    "        optim.zero_grad()\n",
    "        result = model(X)\n",
    "        loss = bceloss(result, y)\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "        print(f\"Loss at epoch number {ep}: {loss.data}\")\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ece12e3a-eef2-438e-bef9-be0636a5b0ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loss at epoch number 1: 0.8681004047393799\n",
      "Loss at epoch number 2: -0.4236203134059906\n",
      "Loss at epoch number 3: -1.2133352756500244\n"
     ]
    }
   ],
   "source": [
    "model = train_model(model, optim, bceloss, X_train, y_train, epochs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "9e28781b-65a5-402d-af1e-05f66e6ccf94",
   "metadata": {},
   "outputs": [],
   "source": [
    "def accuracy(model, X, y):\n",
    "    out = model(X)\n",
    "    correct = torch.abs(y - out) < 0.5\n",
    "    return correct.float().mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "4af589af-029e-403d-836e-2e0957eb37b4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy on plain test_set using a Logistic Regression with a single neural network with a single node. : 0.43565526604652405\n"
     ]
    }
   ],
   "source": [
    "p_accuracy = accuracy(model, X_test, y_test)\n",
    "print(f\"Accuracy on plain test_set using a Logistic Regression with a single neural network with a single node. : {p_accuracy}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e58fe77-0cb0-44c0-8ce4-1d38414715dd",
   "metadata": {},
   "source": [
    "## Encrypted Evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a90f5d00-c20e-4a29-ae35-9149bff873da",
   "metadata": {},
   "source": [
    "Now, we shall look into creating a Logistic Regression model with **plain parameters** at the moment, we can also use encrypted parameters but would try that in later steps. And we will evaluate this model on an encrypted test data set. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c94fe9f0-9426-40ec-af9f-93069d6f47b2",
   "metadata": {},
   "source": [
    "Let's start by creating a pytorch LR model that works for encrypted data,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "6839e153-3d6d-4323-9d55-6f223f2bec35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Author : Tenseal - https://github.com/OpenMined/\n",
    "\n",
    "class EncryptedLR:\n",
    "\n",
    "    def __init__(self, torch_lr):\n",
    "        # TenSEAL processes lists and not torch tensors,\n",
    "        # so we take out the parameters from the PyTorch model\n",
    "        self.weight = torch_lr.lr.weight.data.tolist()[0]\n",
    "        self.bias = torch_lr.lr.bias.data.tolist()\n",
    "\n",
    "    def forward(self, enc_x):\n",
    "        # We don't need to perform sigmoid as this model\n",
    "        # will only be used for evaluation, and the label\n",
    "        # can be deduced without applying sigmoid\n",
    "        enc_out = enc_x.dot(self.weight) + self.bias\n",
    "        return enc_out\n",
    "\n",
    "    def __call__(self, *args, **kwargs):\n",
    "        return self.forward(*args, **kwargs)\n",
    "\n",
    "    def encrypt(self, context):\n",
    "        self.weight = ts.ckks_vector(context, self.weight)\n",
    "        self.bias = ts.ckks_vector(context, self.bias)\n",
    "\n",
    "    def decrypt(self, context):\n",
    "        self.weight = self.weight.decrypt()\n",
    "        self.bias = self.bias.decrypt()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "63ac665a-b8d7-4d95-8651-3f1a9f576632",
   "metadata": {},
   "outputs": [],
   "source": [
    "eelr = EncryptedLR(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c75655ba-d9a1-4037-8667-6071082a9277",
   "metadata": {},
   "source": [
    "Let's create a tenseal context and specify the Homomorphic Encrytion scheme and the define the parameters we want to use. \n",
    "\n",
    "We will choose small and secure parameters that will enable us to perform single multiplication. This should be adequate for evaluating the Logistic Regression model. But for training the model on encrypted data, we would require larger parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "7a09def6-563b-404d-8039-6c9916242fa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define parameters\n",
    "poly_mod_degree = 4096\n",
    "coeff_mod_bit_sizes = [40, 20, 40]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "e968f809-0506-4dad-a490-6781657f1fc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create TenSEALContext\n",
    "ctx_eval = ts.context(ts.SCHEME_TYPE.CKKS, poly_mod_degree, -1, coeff_mod_bit_sizes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "9737a01a-eb1b-408d-aa03-f574c21b529e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Scale of ciphertext to use\n",
    "ctx_eval.global_scale = 2 ** 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "efb0aa47-4941-4c2f-acc9-66a4fad68915",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This key is needed for doing dot-product operations\n",
    "ctx_eval.generate_galois_keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0212fccd-933a-404b-bbb1-0d227b5f15b5",
   "metadata": {},
   "source": [
    "Let's encrypt the test data set before the evaluation process. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "2ec153b5-d970-4bad-a029-3ae0354d9424",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Encryption of the test-set took 1 second(s)\n"
     ]
    }
   ],
   "source": [
    "t_start = time()\n",
    "enc_X_test = [ts.ckks_vector(ctx_eval, x.tolist()) for x in X_test]\n",
    "t_end = time()\n",
    "print(f\"Encryption of the test-set took {int(t_end - t_start)} second(s)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "37559f2e-9876-412b-9395-9a753180583e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# encrypt the model's parameters\n",
    "eelr.encrypt(ctx_eval)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7731c932-18c0-44c3-8a00-65be9ecb4693",
   "metadata": {},
   "source": [
    "We don't need the sigmoid function on the encrypted output layer, however we would need to add this to the encrypted model training. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "1cf98d15-76b9-48ba-88ef-519af34e3b2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def encrypted_evaluation(model, enc_X_test, y_test):\n",
    "    t_start = time()\n",
    "\n",
    "    correct = 0\n",
    "    for enc_X, y in zip(enc_X_test, y_test):\n",
    "        # encrypted evaluation\n",
    "        enc_out = model(enc_X)\n",
    "        # plain comparaison\n",
    "        out = enc_out.decrypt()\n",
    "        out = torch.tensor(out)\n",
    "        out = torch.sigmoid(out)\n",
    "        if torch.abs(out - y) < 0.5:\n",
    "            correct += 1\n",
    "\n",
    "    t_end = time()\n",
    "    print(f\"Evaluated test_set of {len(X_test)} entries in {int(t_end - t_start)} seconds\")\n",
    "    print(f\"Accuracy: {correct}/{len(X_test)} = {correct / len(X_test)}\")\n",
    "    return correct / len(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "b307275e-4a5a-46ae-ba3a-f39e3cbd599d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Evaluated test_set of 847 entries in 2 seconds\n",
      "Accuracy: 133/847 = 0.15702479338842976\n",
      "Difference between plain and encrypted accuracies: 0.27863046526908875\n"
     ]
    }
   ],
   "source": [
    "encrypted_accuracy = encrypted_evaluation(eelr, enc_X_test, y_test)\n",
    "diff_accuracy = p_accuracy - encrypted_accuracy\n",
    "print(f\"Difference between plain and encrypted accuracies: {diff_accuracy}\")\n",
    "if diff_accuracy < 0:\n",
    "    print(\"Looks like we got a better accuracy on the encrypted test-set! The noise was on our end \")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7eca0e22-d897-48c6-a223-5b006985dd50",
   "metadata": {},
   "source": [
    "## Training an Encrypted Logistic Regression model on Encrypted Data\n",
    "\n",
    "Now, we will redefine a PyTorch model that can both forward encrypted data, as well as backpropagate to update the weights and thus train the encrypted logistic regression model on encrypted data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "82b2ffcf-968e-4345-a7ff-bb41c44c3854",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Author : Tenseal - https://github.com/OpenMined/\n",
    "class EncryptedLR:\n",
    "\n",
    "    def __init__(self, torch_lr):\n",
    "        self.weight = torch_lr.lr.weight.data.tolist()[0]\n",
    "        self.bias = torch_lr.lr.bias.data.tolist()\n",
    "        # we accumulate gradients and counts the number of iterations\n",
    "        self._delta_w = 0\n",
    "        self._delta_b = 0\n",
    "        self._count = 0\n",
    "\n",
    "    def forward(self, enc_x):\n",
    "        enc_out = enc_x.dot(self.weight) + self.bias\n",
    "        enc_out = EncryptedLR.sigmoid(enc_out)\n",
    "        return enc_out\n",
    "    \n",
    "    def backward(self, enc_x, enc_out, enc_y):\n",
    "        out_minus_y = (enc_out - enc_y)\n",
    "        self._delta_w += enc_x * out_minus_y\n",
    "        self._delta_b += out_minus_y\n",
    "        self._count += 1\n",
    "\n",
    "    def update_parameters(self):\n",
    "        if self._count == 0:\n",
    "            raise RuntimeError(\"You should at least run one forward iteration\")\n",
    "        # update weights\n",
    "        # We use a small regularization term to keep the output\n",
    "        # of the linear layer in the range of the sigmoid approximation\n",
    "        self.weight -= self._delta_w * (1 / self._count) + self.weight * 0.05\n",
    "        self.bias -= self._delta_b * (1 / self._count)\n",
    "        # reset gradient accumulators and iterations count\n",
    "        self._delta_w = 0\n",
    "        self._delta_b = 0\n",
    "        self._count = 0\n",
    "\n",
    "    @staticmethod\n",
    "    def sigmoid(enc_x):\n",
    "        # We use the polynomial approximation of degree 3\n",
    "        # sigmoid(x) = 0.5 + 0.197 * x - 0.004 * x^3\n",
    "        # from https://eprint.iacr.org/2018/462.pdf\n",
    "        # which fits the function pretty well in the range [-5,5]\n",
    "        return enc_x.polyval([0.5, 0.197, 0, -0.004])\n",
    "\n",
    "    def plain_accuracy(self, x_test, y_test):\n",
    "        # evaluate accuracy of the model on\n",
    "        # the plain (x_test, y_test) dataset\n",
    "        w = torch.tensor(self.weight)\n",
    "        b = torch.tensor(self.bias)\n",
    "        out = torch.sigmoid(x_test.matmul(w) + b).reshape(-1, 1)\n",
    "        correct = torch.abs(y_test - out) < 0.5\n",
    "        return correct.float().mean()\n",
    "    \n",
    "    def encrypt(self, context):\n",
    "        self.weight = ts.ckks_vector(context, self.weight)\n",
    "        self.bias = ts.ckks_vector(context, self.bias)\n",
    "\n",
    "    def decrypt(self):\n",
    "        self.weight = self.weight.decrypt()\n",
    "        self.bias = self.bias.decrypt()\n",
    "\n",
    "    def __call__(self, *args, **kwargs):\n",
    "        return self.forward(*args, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "5e61cb4f-bf1f-410c-8f32-11f98ed93c35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify parameters\n",
    "poly_mod_degree = 8192\n",
    "coeff_mod_bit_sizes = [40, 21, 21, 21, 21, 21, 21, 40]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "5082354a-7ecf-4df5-a26b-748f0e0b2eb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create TenSEALContext\n",
    "ctx_training = ts.context(ts.SCHEME_TYPE.CKKS, poly_mod_degree, -1, coeff_mod_bit_sizes)\n",
    "ctx_training.global_scale = 2 ** 21\n",
    "ctx_training.generate_galois_keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "6bd827be-14f7-4328-b9cc-c78ef870ea8a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Encryption of the training_set took 66 seconds\n"
     ]
    }
   ],
   "source": [
    "t_start = time()\n",
    "enc_X_train = [ts.ckks_vector(ctx_training, x.tolist()) for x in X_train]\n",
    "enc_y_train = [ts.ckks_vector(ctx_training, y.tolist()) for y in y_train]\n",
    "t_end = time()\n",
    "print(f\"Encryption of the training_set took {int(t_end - t_start)} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "906156c9-d6e9-43f1-b159-94d9a9a8121b",
   "metadata": {},
   "source": [
    "We finally reached the last part, which is about training an encrypted logistic regression model on encrypted data! You can see that we decrypt the weights and re-encrypt them again after every epoch, this is necessary since after updating the weights at the end of the epoch, we can no longer use them to perform enough multiplications, so we need to get them back to the initial ciphertext level. In a real scenario, this would translate to sending the weights back to the secret-key holder for decryption and re-encryption. In that case, it will result in just a few Kilobytes of communication per epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "ebb3f962-2d3e-4d45-9de7-8222f638a5f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy at epoch 0 is 0.0\n",
      "Accuracy at epoch #1 is 0.36009445786476135\n",
      "Accuracy at epoch #2 is 0.5938606858253479\n",
      "Accuracy at epoch #3 is 0.5773317813873291\n",
      "\n",
      "Average time per epoch: 350 seconds\n",
      "Final accuracy is 0.5773317813873291\n",
      "Difference between plain and encrypted accuracies: -0.14167651534080505\n",
      "Oh! We got a better accuracy when training on encrypted data!\n"
     ]
    }
   ],
   "source": [
    "eelr = EncryptedLR(LogisticRegression(n_features))\n",
    "accuracy = eelr.plain_accuracy(X_test, y_test)\n",
    "print(f\"Accuracy at epoch 0 is {accuracy}\")\n",
    "\n",
    "times = []\n",
    "for epoch in range(epochs):\n",
    "    eelr.encrypt(ctx_training)\n",
    "    \n",
    "    # if you want to keep an eye on the distribution to make sure\n",
    "    # the function approxiamation is still working fine\n",
    "    # WARNING: this operation is time consuming\n",
    "    # encrypted_out_distribution(eelr, enc_X_train)\n",
    "    \n",
    "    t_start = time()\n",
    "    for enc_X, enc_y in zip(enc_X_train, enc_y_train):\n",
    "        enc_out = eelr.forward(enc_X)\n",
    "        eelr.backward(enc_X, enc_out, enc_y)\n",
    "    eelr.update_parameters()\n",
    "    t_end = time()\n",
    "    times.append(t_end - t_start)\n",
    "    \n",
    "    eelr.decrypt()\n",
    "    accuracy = eelr.plain_accuracy(X_test, y_test)\n",
    "    print(f\"Accuracy at epoch #{epoch + 1} is {accuracy}\")\n",
    "\n",
    "\n",
    "print(f\"\\nAverage time per epoch: {int(sum(times) / len(times))} seconds\")\n",
    "print(f\"Final accuracy is {accuracy}\")\n",
    "\n",
    "diff_accuracy = p_accuracy - accuracy\n",
    "print(f\"Difference between plain and encrypted accuracies: {diff_accuracy}\")\n",
    "if diff_accuracy < 0:\n",
    "    print(\"Oh! We got a better accuracy when training on encrypted data!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c0bd3bd-daec-4a24-84b7-3f1b7d610984",
   "metadata": {},
   "source": [
    "## Conclusion :"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "625f98c9-a6c9-4588-b92b-10f2c73dfdf0",
   "metadata": {},
   "source": [
    "In this notebook, \n",
    "* We explored heart disease data set we found on Kaggle. \n",
    "* We performed EDA and created a simple Logistic Regression model on that. \n",
    "* Then we created a single layer neural network with a single node. \n",
    "* Evaluate this model on an encrypted test data set.\n",
    "* Train an Encrypted Logistic Regression model on Encrypted Data.\n",
    "* Accuracy is low on the training the model on encrypted data and that's because we have lessened the number of epochs to use less resources. But we can get 70-80% accuracy when trained with more epochs. \n",
    "\n",
    "As a part of the future work, we aim to explore more datasets and more machine learning problems that can benefit with the encrypted training and testing. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47808f74-8c3f-43f0-b513-2d614dcd466a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}