main.py

# -*- coding: utf-8 -*-
"""heart-attack-prediction.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/github/talkome/heart-failure-prediction-model/blob/main/heart_attack_prediction.ipynb

# Final Project - Heart Attack Prediction

# Sources

- https://medium.com/pytorch/using-optuna-to-optimize-pytorch-hyperparameters-990607385e36
- https://pytorch.org/
- https://www.pytorchlightning.ai/
- https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html
- https://optuna.org/

# Preparing
"""

import os

for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        print(os.path.join(dirname, filename))

import pandas as pd

pd.options.mode.chained_assignment = None  # default='warn'
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.express as px

from sklearn.utils import shuffle
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing import QuantileTransformer
from sklearn.metrics import confusion_matrix, classification_report
from sklearn.impute import KNNImputer

from sklearn.decomposition import PCA

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim

import optuna
from optuna.trial import TrialState

from sklearn.metrics import f1_score, precision_score, recall_score, accuracy_score, auc, roc_curve, \
    accuracy_score

seed = 42

plt.style.use("seaborn-whitegrid")
plt.rc("figure", autolayout=True)
plt.rc("axes", labelweight="bold", labelsize="large", titleweight="bold", titlesize=14, titlepad=10)

"""## Custom functions definition"""


def get_scores(y, y_pred):
    data = {'Accuracy': np.round(accuracy_score(y, y_pred), 2),
            'Precision': np.round(precision_score(y, y_pred), 2),
            'Recall': np.round(recall_score(y, y_pred), 2),
            'F1': np.round(f1_score(y, y_pred), 2),
            'ROC AUC': np.round(roc_auc(y, y_pred), 2)}
    scores_df = pd.Series(data).to_frame('scores')
    return scores_df


def conf_matrix(y, y_pred):
    fig, ax = plt.subplots(figsize=(3.5, 3.5))
    labels = ['No', 'Yes']
    ax = sns.heatmap(confusion_matrix(y, y_pred), annot=True, cmap="Blues", fmt='g', cbar=False, annot_kws={"size": 25})
    plt.title('Heart Failure?', fontsize=20)
    ax.xaxis.set_ticklabels(labels, fontsize=17)
    ax.yaxis.set_ticklabels(labels, fontsize=17)
    ax.set_ylabel('Test')
    ax.set_xlabel('Predicted')


#
#
# """## Dataset Loading"""
#
df = pd.read_csv('heart.csv', skipinitialspace=True)

df.head()

df.info()

"""# Data Analysis

Are there null values?
"""

df.isnull().sum()

"""There are no null values.

Are there duplicate values?
"""

df[df.duplicated()]

"""There are no duplicate values.

To simplify the following analysis, we will define a function to plot numerical features by plotly.
"""

num_cols = ['Age', 'RestingBP', 'Cholesterol', 'FastingBS', 'MaxHR', 'Oldpeak']


def num_plot(df, col):
    fig = px.histogram(df, x=col, color="HeartDisease",
                       marginal="box")
    fig.update_layout(height=400, width=500, showlegend=True)
    fig.update_traces(marker_line_width=1, marker_line_color="black")
    fig.show()


for col in num_cols:
    num_plot(df, col)

"""There are outliers values for Colesterol and restingBP=0"""

df[df['RestingBP'] == 0]

"""We will drop this value!"""

df = df[(df['RestingBP'] > 0)]

"""What about 0 values of Cholesterol?"""

df[df['Cholesterol'] == 0]

"""There are 171 patients with a cholesterol value = 0, which is not possible."""

num_plot(df, 'Cholesterol')

print('Mean', df.Cholesterol.mean())
print('Median', df.Cholesterol.median())
print('Mode', df.Cholesterol.mode()[0])

"""The most frequent value (mode) is cholesterol = 0."""

df_no_chol = df[df['Cholesterol'] == 0]

for col in num_cols:
    num_plot(df_no_chol, col)

"""It looks like there are no common features among these patients with missing cholesterol value.<br>
We will impute the missing values of Cholesterol after splitting the dataset into training and test sets. For now, we label these values as 'NaN'.
"""

df['Cholesterol'] = df['Cholesterol'].replace({0: np.nan})

"""# Categorical Features Encoding

Depending on the number of different possible values for each categorical variable, we will choose if encoding the variable by label encoding or one hot encoding (OHE).

## Gender

Sex has only two classes, so we can encode it by 0 and 1.
"""

df.Sex = df.Sex.replace({'M': 0, 'F': 1})

"""## ChestPainType"""

df.ChestPainType.value_counts()

"""There are only 3 classes, we can encode it by OHE

## ExerciseAngina
"""

df.ExerciseAngina.value_counts()

"""There are two classes, we will encode it by 0 and 1."""

df.ExerciseAngina = df.ExerciseAngina.replace({'N': 0, 'Y': 1})

"""## ST_Slope"""

df.ST_Slope.value_counts()

"""There are only 3 classes, we can encode it by OHE

## RestingECG
"""

df.RestingECG.value_counts()

"""There are only 3 classes, we can encode it by OHE

We will encode the variables by OHE using the convenient get_dummies function from pandas library.
"""

encoded_df = pd.get_dummies(df, drop_first=True)

"""# Data preprocessing for Machine Learning modeling"""

fig = px.histogram(df, x="HeartDisease", color="HeartDisease")
fig.update_layout(height=400, width=500, showlegend=True)
fig.update_traces(marker_line_width=1, marker_line_color="black")
fig.show()

"""Good news, the target variable looks balanced."""

X = encoded_df.drop('HeartDisease', axis=1)

y = encoded_df['HeartDisease']

"""## Train - Test split"""

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.15, random_state=seed, stratify=y)

"""## Missing Cholesterol values imputation

After extracting the training set, we can study how to impute missing values of the cholesterol column. We perform the imputation only on the training set to avoid any data leakage: the best value for imputation obtained on the training set will be used also on the test set.
"""

fig, ax = plt.subplots(2, 1, sharex=True, figsize=(8, 5), gridspec_kw={"height_ratios": (.2, .8)})
ax[0].set_title('Cholesterol distribution', fontsize=18)
sns.boxplot(x='Cholesterol', data=X_train, ax=ax[0])
ax[0].set(yticks=[])
sns.histplot(x='Cholesterol', data=X_train, ax=ax[1])
ax[1].set_xlabel(col, fontsize=16)
plt.axvline(X_train['Cholesterol'].mean(), color='darkgreen', linewidth=2.2,
            label='mean=' + str(np.round(X_train['Cholesterol'].mean(), 1)))
plt.axvline(X_train['Cholesterol'].median(), color='red', linewidth=2.2,
            label='median=' + str(np.round(X_train['Cholesterol'].median(), 1)))
plt.axvline(X_train['Cholesterol'].mode()[0], color='purple', linewidth=2.2,
            label='mode=' + str(X_train['Cholesterol'].mode()[0]))
plt.legend(bbox_to_anchor=(1, 1.03), ncol=1, fontsize=17, fancybox=True, shadow=True, frameon=True)
plt.tight_layout()
plt.show()

"""By visual inspection, a good value for imputing missing cholesterol values could be 240."""

chol = 240

X_train['Cholesterol'] = X_train['Cholesterol'].fillna(chol)

X_test['Cholesterol'] = X_test['Cholesterol'].fillna(chol)

scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

"""# Dimensionality reduction by PCA

We have 15 features/columns in the dataset, is it possible to remove some of them and still keep a high explainability?
"""

pca = PCA()
pca.fit_transform(X_train);

cum_sum = np.cumsum(pca.explained_variance_ratio_) * 100
comp = [n for n in range(len(cum_sum))]

plt.figure(figsize=(5, 5))
plt.plot(comp, cum_sum, marker='.')
plt.xlabel('PCA Components')
plt.ylabel('Cumulative Explained Variance (%)')
plt.title('PCA')
plt.show()

"""It looks like the PCA curve does not have any clear elbow. We should not drop any of the features.

## Train - Validation split
"""

X_train, X_valid, y_train, y_valid = train_test_split(X_train, y_train, test_size=0.2, stratify=y_train,
                                                      random_state=seed)

"""## Create Custom Dataset class"""

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(DEVICE)


class CustomDataset:
    def __init__(self, X_data, y_data, device=DEVICE):
        self.X_data = X_data
        self.y_data = y_data

    def __len__(self):
        return len(self.X_data)

    def __getitem__(self, index):
        return self.X_data[index], self.y_data[index]


train_data = CustomDataset(torch.FloatTensor(X_train), torch.FloatTensor(y_train.values))

val_data = CustomDataset(torch.FloatTensor(X_valid), torch.FloatTensor(y_valid.values))

test_data = CustomDataset(torch.FloatTensor(X_test), torch.FloatTensor(y_test.values))

"""## Parameters definition"""

BATCHSIZE = 16

"""## Dataloader definition"""

train_loader = DataLoader(dataset=train_data, batch_size=BATCHSIZE)
valid_loader = DataLoader(dataset=val_data, batch_size=1)
test_loader = DataLoader(dataset=test_data, batch_size=1)

"""# Hyperoptimization by Optuna

# Neural Network hypermodel

In the following we will define a **n-layers Neural Network hypermodel**, where n is the number of hidden layers chosen by OPTUNA. The idea of a hypermodel is that some hyperparameters are defined in a range of values, where the optmization algorithm (OPTUNA) will search to find the 'best' ones by optimizing a certain metric, for example the validation accuracy of the model. <br>
In this work, the hyperparameters that will be optimized are:
- number of layers
- number of neurons for each layer
- learning rate of Adam optimizer
"""


def define_model(trial):
    n_layers = trial.suggest_int("n_layers", 1, 2)
    layers = []

    in_features = 15
    for i in range(n_layers):
        out_features = trial.suggest_int("n_units_{}".format(i), 8, 25)
        layers.append(nn.Linear(in_features, out_features))
        layers.append(nn.ReLU())
        p = trial.suggest_uniform("dropout_{}".format(i), 0.2, 0.5)
        layers.append(nn.Dropout(p))
        in_features = out_features
    layers.append(nn.Linear(out_features, 1))

    return nn.Sequential(*layers)


"""Then, we need to define an objective function, including the train and evaluation of our model.

# OPTUNA
"""

EPOCHS = 40


def objective(trial):
    # call the define_model method
    model = define_model(trial).to(DEVICE)

    # Optimizer and loss definition
    lr = trial.suggest_float("lr", 5e-4, 1e-2, log=True)
    optimizer = getattr(optim, 'Adam')(model.parameters(), lr=lr)
    criterion = nn.BCEWithLogitsLoss()
    # Using the logit binary CE, we include the sigmoid function in the prediction output during the loss calculation

    train_acc = []
    train_loss = []

    valid_acc = []
    valid_loss = []

    total_step = len(train_loader)
    total_step_val = len(valid_loader)

    for epoch in range(EPOCHS):

        running_loss = 0
        correct = 0
        total = 0

        # TRAINING

        model.train()

        for batch_idx, (X_train_batch, y_train_batch) in enumerate(train_loader):
            X_train_batch, y_train_batch = X_train_batch.to(DEVICE), y_train_batch.to(DEVICE)
            optimizer.zero_grad()
            output = model(X_train_batch)
            y_pred = torch.round(torch.sigmoid(output))
            # LOSS
            loss = criterion(output, y_train_batch.unsqueeze(1))
            loss.backward()
            optimizer.step()
            running_loss += loss.item()  # sum all batch losses
            # ACCURACY
            correct += torch.sum(y_pred == y_train_batch.unsqueeze(1)).item()
            total += y_train_batch.size(0)
        train_acc.append(100 * correct / total)
        train_loss.append(
            running_loss / total_step)  # get average loss among all batches dividing total loss by the number of batches

        # VALIDATION
        correct_v = 0
        total_v = 0
        batch_loss = 0
        with torch.no_grad():
            model.eval()
            for batch_idx, (X_valid_batch, y_valid_batch) in enumerate(valid_loader):
                X_valid_batch, y_valid_batch = X_valid_batch.to(DEVICE), y_valid_batch.to(DEVICE)
                # PREDICTION
                output = model(X_valid_batch)
                y_pred = torch.round(torch.sigmoid(output))
                # LOSS
                loss_v = criterion(output, y_valid_batch.unsqueeze(1))
                batch_loss += loss_v.item()
                # ACCURACY
                correct_v += torch.sum(y_pred == y_valid_batch.unsqueeze(1)).item()
                total_v += y_valid_batch.size(0)
            valid_acc.append(100 * correct_v / total_v)
            valid_loss.append(batch_loss / total_step_val)

        trial.report(np.mean(valid_loss), epoch)

        # Handle pruning based on the intermediate value
        if trial.should_prune():
            raise optuna.exceptions.TrialPruned()

    return np.mean(valid_loss)


"""Finally, we can start the optimization with Optuna."""

study = optuna.create_study(direction="minimize")
study.optimize(objective, n_trials=100)

pruned_trials = study.get_trials(deepcopy=False, states=[TrialState.PRUNED])
complete_trials = study.get_trials(deepcopy=False, states=[TrialState.COMPLETE])

print("Study statistics: ")
print("  Number of finished trials: ", len(study.trials))
print("  Number of pruned trials: ", len(pruned_trials))
print("  Number of complete trials: ", len(complete_trials))

print("Best trial:")
trial = study.best_trial

print("  Value: ", trial.value)

print("  Params: ")

params = []
#
for key, value in trial.params.items():
    params.append(value)
    print("    {}: {}".format(key, value))

params

"""We can extract the best parameters from the list:"""

n_layers = params[0]
#
units_1 = params[1]
dropout_1 = np.round(params[2], 5)
#
lr = np.round(params[3], 8)

"""Then, we define a new Neural Network by Pytorch, where the hyperparameters will have the optimized values obtaiend with optuna.

# Pytorch Neural Network modeling
"""


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layer_1 = nn.Linear(X_train.shape[1], units_1)
        self.layer_out = nn.Linear(units_1, 1)
        self.dropout1 = nn.Dropout(p=dropout_1)

    def forward(self, inputs):
        x = F.relu(self.layer_1(inputs))
        x = self.dropout1(x)
        x = self.layer_out(x)

        return x


model = Net()
model.to(DEVICE)
print(model)

criterion = nn.BCEWithLogitsLoss()
optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.0001)

"""We set the number of epochs:"""

EPOCHS = 100

"""# Pytorch Neural Network training

Finally we can start the training of the Neural Network with the optimized hyperparameters
"""

# Model Training

early_stopping_patience = 15
early_stopping_counter = 0

train_acc = []
train_loss = []

valid_acc = []
valid_loss = []

total_step = len(train_loader)
total_step_val = len(valid_loader)

valid_loss_min = np.inf

for epoch in range(EPOCHS):

    running_loss = 0
    correct = 0
    total = 0

    # TRAINING

    model.train()

    for batch_idx, (X_train_batch, y_train_batch) in enumerate(train_loader):
        X_train_batch, y_train_batch = X_train_batch.to(DEVICE), y_train_batch.to(DEVICE)
        optimizer.zero_grad()
        output = model(X_train_batch)
        y_pred = torch.round(torch.sigmoid(output))
        # LOSS
        loss = criterion(output, y_train_batch.unsqueeze(1))
        loss.backward()
        optimizer.step()
        running_loss += loss.item()  # sum loss for every batch
        # ACCURACY
        correct += torch.sum(y_pred == y_train_batch.unsqueeze(1)).item()
        total += y_train_batch.size(0)
    train_acc.append(100 * correct / total)  # calculate accuracy among all entries in the batches
    train_loss.append(
        running_loss / total_step)  # get average loss among all batches dividing total loss by the number of batches

    # VALIDATION
    correct_v = 0
    total_v = 0
    batch_loss = 0
    with torch.no_grad():
        model.eval()
        for batch_idx, (X_valid_batch, y_valid_batch) in enumerate(valid_loader):
            X_valid_batch, y_valid_batch = X_valid_batch.to(DEVICE), y_valid_batch.to(DEVICE)
            # PREDICTION
            output = model(X_valid_batch)
            y_pred = torch.round(torch.sigmoid(output))
            # LOSS
            loss_v = criterion(output, y_valid_batch.unsqueeze(1))
            batch_loss += loss_v.item()
            # ACCURACY
            correct_v += torch.sum(y_pred == y_valid_batch.unsqueeze(1)).item()
            total_v += y_valid_batch.size(0)
        valid_acc.append(100 * correct_v / total_v)
        valid_loss.append(batch_loss / total_step_val)

    if np.mean(valid_loss) <= valid_loss_min:
        torch.save(model.state_dict(), './state_dict.pt')
        print(
            f'Epoch {epoch + 0:01}: Validation loss decreased ({valid_loss_min:.6f} --> {np.mean(valid_loss):.6f}).  Saving model ...')
        valid_loss_min = np.mean(valid_loss)
        early_stopping_counter = 0  # reset counter if validation loss decreases
    else:
        print(f'Epoch {epoch + 0:01}: Validation loss did not decrease')
        early_stopping_counter += 1

    if early_stopping_counter > early_stopping_patience:
        print('Early stopped at epoch :', epoch)
        break

    print(
        f'\t Train_Loss: {np.mean(train_loss):.4f} Train_Acc: {(100 * correct / total):.3f} Val_Loss: {np.mean(valid_loss):.4f}  BEST VAL Loss: {valid_loss_min:.4f}  Val_Acc: {(100 * correct_v / total_v):.3f}\n')

"""Now we can get the predictions by looping over the test loader."""

y_pred_prob_list = []
y_pred_list = []

# Loading the best model
model.load_state_dict(torch.load('./state_dict.pt'))

with torch.no_grad():
    model.eval()
    ## Added feature
    ##data = torch.randn(58, 0, 15, 136, 164, 0, 15, 99, 15, 2, 15)
    data = torch.randn(55, 0, 0, 120, 270, 0, 0, 140, 0, 0, 15)
    print(data)
    output = model(data)
    print(output)
    prediction = torch.argmax(output)
    print(prediction)
#     for batch_idx, (X_test_batch, y_test_batch) in enumerate(test_loader):
#         X_test_batch = X_test_batch.to(DEVICE)
#         # PREDICTION
#         output = model(X_test_batch)
#         y_pred_prob = torch.sigmoid(output)
#         y_pred_prob_list.append(y_pred_prob.cpu().numpy())
#         y_pred = torch.round(y_pred_prob)
#         y_pred_list.append(y_pred.cpu().numpy())
#
# y_pred_prob_list = [a.squeeze().tolist() for a in y_pred_prob_list]
# y_pred_list = [a.squeeze().tolist() for a in y_pred_list]

"""# PyTorch Results Summary

## Classification Report
"""

# print(classification_report(y_test, y_pred_list))
#
#
# """## Confusion Matrix"""
#
# conf_matrix(y_test, y_pred_list)
#
# """## PyTorch NN ROC curve"""
#
# plt.figure(figsize=(5.5, 4))
#
# fpr, tpr, _ = roc_curve(y_test, y_pred_prob_list)
# roc_auc = auc(fpr, tpr)
# plt.plot(fpr, tpr, 'b', label='AUC = %0.2f' % roc_auc)
# plt.plot([0, 1], [0, 1], 'r--')
# plt.title('ROC curve', fontsize=25)
# plt.ylabel('True Positive Rate', fontsize=18)
# plt.xlabel('False Positive Rate', fontsize=18)
# plt.legend(loc='lower right', fontsize=24, fancybox=True, shadow=True, frameon=True, handlelength=0)
# plt.show()

"""**The results of the hyperoptimized Pytorch Neural Network are satisfying, since all the scores are around 90%.**<br>
**In particular, the best Recall score I could get is 92%, which in medical applications is the best score to look at, since having false negatives is a crime**
"""


def predict_patient():
    l = [58, 0, "ATA", 136, 164, 0, "ST", 99, "Y", 2, "Flat", 1]