PlotUtils.py 1.3 KB

123456789101112131415161718192021222324252627282930313233
  1. import matplotlib.pyplot as plt
  2. from sklearn.metrics import roc_curve, auc
  3. def plot_roc_curve(test_labels: list, test_df: list, title: str, figsize=(8, 8), savefile = None, show: bool = True):
  4. fpr, tpr, thresholds = roc_curve(test_labels, test_df)
  5. auc_score = auc(fpr, tpr)
  6. plt.figure(figsize=figsize)
  7. plt.plot(fpr, tpr, lw=1)
  8. plt.fill_between(fpr, tpr, label=f"AUC = {auc_score:.4f}", alpha=0.5)
  9. plt.plot([0, 1], [0, 1], color="gray", linestyle="dotted")
  10. plt.xlim([0.0, 1.0])
  11. plt.ylim([0.0, 1.0])
  12. plt.xlabel("FPR")
  13. plt.ylabel("TPR")
  14. plt.title(f"{title}")
  15. plt.legend(loc="lower right")
  16. if savefile is not None:
  17. plt.savefig(f"{savefile}.png", bbox_inches="tight")
  18. plt.savefig(f"{savefile}.pdf", bbox_inches="tight")
  19. if show:
  20. plt.show()
  21. return fpr, tpr, thresholds, auc_score
  22. def get_percentiles(fpr, tpr, thresholds, percentiles=[0.9, 0.95, 0.98, 0.99]):
  23. tnrs = []
  24. for percentile in percentiles:
  25. for i, tp in enumerate(tpr):
  26. if tp >= percentile:
  27. tnrs.append(1 - fpr[i]) # append tnr
  28. print(f"{percentile} percentile : TPR = {tp:.4f}, FPR = {fpr[i]:.4f} <-> TNR = {(1 - fpr[i]):.4f} @ thresh {thresholds[i]}")
  29. break
  30. return tnrs