main.py

# Importing modules
import torch
import numpy as np
import pandas as pd
from utils.dataDownload import download_freeSolv
from sklearn.model_selection import train_test_split
from utils.dataloader import loadData
from src.model import GCN
from src.train import TrainGCN
from src.test import TestGCN
from sklearn.metrics import classification_report
from sklearn.metrics import confusion_matrix
from sklearn.metrics import ConfusionMatrixDisplay
from sklearn.metrics import classification_report
import matplotlib.pyplot as plt

# Download FreeSolv data from moleculenet
download_freeSolv()

# Loading dataset
data = pd.read_csv("data/data.csv")

# Smiles
X = data["smiles"].to_numpy()

# Train/Test Split
X_train, X_test = train_test_split(X, test_size=0.2, random_state=42)

# Train loader
train_loader = loadData(X_train)

# Test loader
test_loader =  loadData(X_test, augment=False)

# GCN model
model = GCN(num_features=7, hidden_dim=64, num_classes=3)

# Loading to device
device = torch.device("cpu")
model = model.to(device)

# Model training
model, plot_data = TrainGCN(model, train_loader, epochs=50)
print("Model Training Complete ...")

# Saving model
torch.save(model, "results/model.pt")
print("Saving model: `results/model.pt`")

# Loss/Accuracy curve
plt.plot(plot_data["epoch"], plot_data["loss"], color="blue", label="Loss")
plt.plot(plot_data["epoch"], plot_data["accuracy"], color="orange", label="Accuracy")
plt.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
plt.savefig("results/loss_accuracy_curve.png")
plt.close()

# Model testing
y_test, y_pred = TestGCN(model, test_loader)
print("Model Testing Complete ...")

# Relabelling target variables
temp = {0:"C", 1:"N", 2:"O"}
y_test = [temp[i] for i in y_test]
y_pred = [temp[i] for i in y_pred]

# Classification report
report = open("results/report.txt", "w")
report.write(classification_report(y_test, y_pred))
report.close()

# Confusion_matrix
cmat = confusion_matrix(y_test, y_pred, normalize="true", labels=["C", "N", "O"])
disp_cmat = ConfusionMatrixDisplay(confusion_matrix=cmat, display_labels=["C", "N", "O"])
disp_cmat.plot()
plt.savefig("results/confusion_matrix.png")
plt.close()

print("Results:\n`results/loss_accuracy_curve.png`\n`results/report.txt`\n`results/confusion_matrix.png`")

# Usage:
#	$ python main.py