<repo_name>boorac/machine-learning<file_sep>/l4/mymlutils.py
from sklearn.metrics import zero_one_loss
import matplotlib.pyplot as plt

def calculate_losses(clf, X_train, X_test, y_train, y_test):
    clf.fit(X_train, y_train)

    predicted_train = clf.predict(X_train)
    predicted_test = clf.predict(X_test)

    loss_train = zero_one_loss(y_train, predicted_train)
    loss_test = zero_one_loss(y_test, predicted_test)

    return loss_train, loss_test

def plot_losses_against_hyperparameter(clfs, hyperparameter_range, hyperparameter_name, X_train, X_test, y_train, y_test, return_optimal=False):
    losses_train = []
    losses_test = []
    figures = []

    optimal_loss = [hyperparameter_range[0], float("inf")]

    for clf, hyperparameter in zip(clfs, hyperparameter_range):
        loss_train, loss_test = calculate_losses(clf, X_train, X_test, y_train, y_test)
        losses_train.append(loss_train)
        losses_test.append(loss_test)

        if return_optimal:
            if(optimal_loss[1] > loss_test):
                optimal_loss[0] = hyperparameter
                optimal_loss[1] = loss_test


    fig_loss_test, = plt.plot(hyperparameter_range, losses_test)
    fig_loss_train, = plt.plot(hyperparameter_range, losses_train)

    plt.legend([fig_loss_test, fig_loss_train], ["test loss", "train_loss"])
    plt.ylabel("E")
    plt.xlabel(hyperparameter_name)
    plt.show()

    if return_optimal:
        return optimal_loss[0], optimal_loss[1]
