Search code examples
pythonmatplotlibscikit-learnprecision-recall

Plotting Threshold (precision_recall curve) matplotlib/sklearn.metrics

I am trying to plot the thresholds for my precision/recall curve. I am just using the MNSIT data, with the example from the book Hands On Machine Learning with scikit-learn, keras, and TensorFlow. Trying to train the model to detect the image of 5's. I don't know how much of the code you need to see. I have made my confusion matrix for the training set and have calculated the precision and recall values, along with the thresholds. I have plotted the pre/rec curve and the example in the book says to add axis label, ledged, grid and highlight the thresholds but the code cuts off in the book where I placed an asterisk below. I was able to figure out all but how to get the thresholds to show up on the plot. I have included a picture of what the graph in the book looks like vs what I have. This is what the book shows:

enter image description here vs my graph:

enter image description here

I can't get that red dotline with two threshold points to show up. Does anyone have any idea how I would do this? Here is my code below:

from sklearn.metrics import precision_recall_curve

precisions, recalls, thresholds = precision_recall_curve(y_train_5, y_scores)

def plot_precision_recall_vs_thresholds(precisions, recalls, thresholds):
    plt.plot(thresholds, precisions[:-1], "b--", label="Precision")
    plt.plot(thresholds, recalls[:-1], "g--", label="Recall")
    plt.xlabel("Threshold")
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0.)
    plt.grid(b=True, which="both", axis="both", color='gray', linestyle='-', linewidth=1)

plot_precision_recall_vs_thresholds(precisions, recalls, thresholds)
plt.show()

I know there's a decent amount of questions on here about this with sklearn but none seem to cover getting that red line to show up. I would greatly appreciate the help!

Solution

  • You can use the following code for plotting horizontal and vertical lines:

    plt.axhline(y_value, c='r', ls=':')
plt.axvline(x_value, c='r', ls=':')