auc.py

"""
Calculate AUC of a sample
Reference:
http://binf.gmu.edu/mmasso/ROC101.pdf
"""
import matplotlib.pyplot as plt

# each element is a tuple(positive, score)
sample = [
  (True, 0.63),
  (True, 0.82),
  (True, 0.53),
  (False, 0.83),
  (False, 0.32),
  (True, 0.89),
  (True, 0.98),
  (False, 0.12),
  (False, 0.13),
  (False, 0.34),
  (True, 0.89),
  (True, 0.63),
  (True, 0.98),
  (False, 0.42),
  (True, 0.63),
  (False, 0.89),
  (False, 0.63),
  (False, 0.62),
  (False, 0.76),
  (False, 0.72),
  (False, 0.78),
  (False, 0.57),
  (False, 0.58),
  (False, 0.63),
  (False, 0.63)
]

FMIN = -float('inf')

def roc(sample):
  sample = sorted(sample, key=lambda x: x[1], reverse=True)
  N = float([item[0] for item in sample].count(False))
  P = len(sample) - N
  fp = tp = 0
  prevScore = FMIN
  result = []
  for item in sample:
    positive, score = item
    if score != prevScore:
      result.append((fp / N, tp / P))
      prevScore = score
    if positive:
      tp = tp + 1
    else:
      fp = fp + 1
    result.append((fp / N, tp / P))
  return result

def auc(sample):
  def update_area(x1, y1, x2, y2):
    return area + abs(x1 - x2) * (y1 + y2) / 2.0

  sample = sorted(sample, key=lambda x: x[1], reverse=True)
  N = float([item[0] for item in sample].count(False))
  P = len(sample) - N
  prevfp = prevtp = fp = tp = area = 0
  prevScore = FMIN
  for item in sample:
    positive, score = item
    if score != prevScore:
      area = update_area(fp, tp, prevfp, prevtp)
      prevScore = score
      prevtp = tp
      prevfp = fp
    if positive:
      tp = tp + 1
    else:
      fp = fp + 1
  # There's a bug in the algorithm presented in the paper
  area = update_area(N, P, prevfp, prevtp) 
  area = area / (P * N) # scale from P * N onto the unit square
  return area

def plot(points):
  x = [p[0] for p in points]
  y = [p[1] for p in points]
  plt.plot(x, y, 'ro')
  plt.plot(x, y, 'b')
  plt.show()

if __name__ == '__main__':
  r = roc(sample)
  a = auc(sample)
  print a
  plot(r)