nupic_anomaly_output.py

# ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2013, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------
"""
Provides two classes with the same signature for writing data out of NuPIC
models.
(This is a component of the One Hot Gym Anomaly Tutorial.)
"""
import csv
from collections import deque
from abc import ABCMeta, abstractmethod
from nupic.algorithms import anomaly_likelihood
# Try to import matplotlib, but we don't have to.
try:
  import matplotlib
  matplotlib.use('TKAgg')
  import matplotlib.pyplot as plt
  import matplotlib.gridspec as gridspec
  from matplotlib.dates import date2num, DateFormatter
except ImportError:
  pass

WINDOW = 300
HIGHLIGHT_ALPHA = 0.3
ANOMALY_HIGHLIGHT_COLOR = 'red'
WEEKEND_HIGHLIGHT_COLOR = 'yellow'
ANOMALY_THRESHOLD = 0.9


class NuPICOutput(object):

  __metaclass__ = ABCMeta


  def __init__(self, name):
    self.name = name
    self.anomalyLikelihoodHelper = anomaly_likelihood.AnomalyLikelihood()


  @abstractmethod
  def write(self, timestamp, value, predicted, anomalyScore):
    pass


  @abstractmethod
  def close(self):
    pass




class NuPICFileOutput(NuPICOutput):


  def __init__(self, *args, **kwargs):
    super(NuPICFileOutput, self).__init__(*args, **kwargs)
    # Prevent from creating local File for output
    '''
    self.outputFiles = []
    self.outputWriters = []
    self.lineCount = 0
    headerRow = [
      'timestamp', 'heart_beat', 'prediction',
      'anomaly_score', 'anomaly_likelihood'
    ]
    outputFileName = "%s_out.csv" % self.name
    print "Preparing to output %s data to %s" % (self.name, outputFileName)
    self.outputFile = open(outputFileName, "w")
    self.outputWriter = csv.writer(self.outputFile)
    self.outputWriter.writerow(headerRow)
    '''



  def write(self, timestamp, value, predicted, anomalyScore):
    if timestamp is not None:
      anomalyLikelihood = self.anomalyLikelihoodHelper.anomalyProbability(
        value, anomalyScore, timestamp
      )
      outputRow = [timestamp, value, predicted, anomalyScore, anomalyLikelihood]
      # self.outputWriter.writerow(outputRow) # Prevent from creating local File for output
      # self.lineCount += 1
      return anomalyLikelihood



  def close(self):
    self.outputFile.close()
    print "Done. Wrote %i data lines to %s." % (self.lineCount, self.name)



def extractWeekendHighlights(dates):
  weekendsOut = []
  weekendSearch = [5, 6]
  weekendStart = None
  for i, date in enumerate(dates):
    if date.weekday() in weekendSearch:
      if weekendStart is None:
        # Mark start of weekend
        weekendStart = i
    else:
      if weekendStart is not None:
        # Mark end of weekend
        weekendsOut.append((
          weekendStart, i, WEEKEND_HIGHLIGHT_COLOR, HIGHLIGHT_ALPHA
        ))
        weekendStart = None

  # Cap it off if we're still in the middle of a weekend
  if weekendStart is not None:
    weekendsOut.append((
      weekendStart, len(dates)-1, WEEKEND_HIGHLIGHT_COLOR, HIGHLIGHT_ALPHA
    ))

  return weekendsOut



def extractAnomalyIndices(anomalyLikelihood):
  anomaliesOut = []
  anomalyStart = None
  for i, likelihood in enumerate(anomalyLikelihood):
    if likelihood >= ANOMALY_THRESHOLD:
      if anomalyStart is None:
        # Mark start of anomaly
        anomalyStart = i
    else:
      if anomalyStart is not None:
        # Mark end of anomaly
        anomaliesOut.append((
          anomalyStart, i, ANOMALY_HIGHLIGHT_COLOR, HIGHLIGHT_ALPHA
        ))
        anomalyStart = None

  # Cap it off if we're still in the middle of an anomaly
  if anomalyStart is not None:
    anomaliesOut.append((
      anomalyStart, len(anomalyLikelihood)-1,
      ANOMALY_HIGHLIGHT_COLOR, HIGHLIGHT_ALPHA
    ))

  return anomaliesOut



class NuPICPlotOutput(NuPICOutput):


  def __init__(self, *args, **kwargs):
    super(NuPICPlotOutput, self).__init__(*args, **kwargs)
    # Turn matplotlib interactive mode on.
    plt.ion()
    self.dates = []
    self.convertedDates = []
    self.value = []
    self.allValues = []
    self.predicted = []
    self.anomalyScore = []
    self.anomalyLikelihood = []
    self.actualLine = None
    self.predictedLine = None
    self.anomalyScoreLine = None
    self.anomalyLikelihoodLine = None
    self.linesInitialized = False
    self._chartHighlights = []
    fig = plt.figure(figsize=(16, 10))
    gs = gridspec.GridSpec(2, 1, height_ratios=[3,  1])

    self._mainGraph = fig.add_subplot(gs[0, 0])
    plt.title(self.name)
    plt.ylabel('KW Energy Consumption')
    plt.xlabel('Date')

    self._anomalyGraph = fig.add_subplot(gs[1])

    plt.ylabel('Percentage')
    plt.xlabel('Date')

    # Maximizes window
    mng = plt.get_current_fig_manager()
    mng.resize(*mng.window.maxsize())

    plt.tight_layout()



  def initializeLines(self, timestamp):
    print "initializing %s" % self.name
    anomalyRange = (0.0, 1.0)
    self.dates = deque([timestamp] * WINDOW, maxlen=WINDOW)
    self.convertedDates = deque(
      [date2num(date) for date in self.dates], maxlen=WINDOW
    )
    self.value = deque([0.0] * WINDOW, maxlen=WINDOW)
    self.predicted = deque([0.0] * WINDOW, maxlen=WINDOW)
    self.anomalyScore = deque([0.0] * WINDOW, maxlen=WINDOW)
    self.anomalyLikelihood = deque([0.0] * WINDOW, maxlen=WINDOW)

    actualPlot, = self._mainGraph.plot(self.dates, self.value)
    self.actualLine = actualPlot
    predictedPlot, = self._mainGraph.plot(self.dates, self.predicted)
    self.predictedLine = predictedPlot
    self._mainGraph.legend(tuple(['actual', 'predicted']), loc=3)

    anomalyScorePlot, = self._anomalyGraph.plot(
      self.dates, self.anomalyScore, 'm'
    )
    anomalyScorePlot.axes.set_ylim(anomalyRange)

    self.anomalyScoreLine = anomalyScorePlot
    anomalyLikelihoodPlot, = self._anomalyGraph.plot(
      self.dates, self.anomalyScore, 'r'
    )
    anomalyLikelihoodPlot.axes.set_ylim(anomalyRange)
    self.anomalyLikelihoodLine = anomalyLikelihoodPlot
    self._anomalyGraph.legend(
      tuple(['anomaly score', 'anomaly likelihood']), loc=3
    )

    dateFormatter = DateFormatter('%m/%d %H:%M')
    self._mainGraph.xaxis.set_major_formatter(dateFormatter)
    self._anomalyGraph.xaxis.set_major_formatter(dateFormatter)

    self._mainGraph.relim()
    self._mainGraph.autoscale_view(True, True, True)

    self.linesInitialized = True



  def highlightChart(self, highlights, chart):
    for highlight in highlights:
      # Each highlight contains [start-index, stop-index, color, alpha]
      self._chartHighlights.append(chart.axvspan(
        self.convertedDates[highlight[0]], self.convertedDates[highlight[1]],
        color=highlight[2], alpha=highlight[3]
      ))



  def write(self, timestamp, value, predicted, anomalyScore):

    # We need the first timestamp to initialize the lines at the right X value,
    # so do that check first.
    if not self.linesInitialized:
      self.initializeLines(timestamp)

    anomalyLikelihood = self.anomalyLikelihoodHelper.anomalyProbability(
      value, anomalyScore, timestamp
    )

    self.dates.append(timestamp)
    self.convertedDates.append(date2num(timestamp))
    self.value.append(value)
    self.allValues.append(value)
    self.predicted.append(predicted)
    self.anomalyScore.append(anomalyScore)
    self.anomalyLikelihood.append(anomalyLikelihood)

    # Update main chart data
    self.actualLine.set_xdata(self.convertedDates)
    self.actualLine.set_ydata(self.value)
    self.predictedLine.set_xdata(self.convertedDates)
    self.predictedLine.set_ydata(self.predicted)
    # Update anomaly chart data
    self.anomalyScoreLine.set_xdata(self.convertedDates)
    self.anomalyScoreLine.set_ydata(self.anomalyScore)
    self.anomalyLikelihoodLine.set_xdata(self.convertedDates)
    self.anomalyLikelihoodLine.set_ydata(self.anomalyLikelihood)

    # Remove previous highlighted regions
    for poly in self._chartHighlights:
      poly.remove()
    self._chartHighlights = []

    weekends = extractWeekendHighlights(self.dates)
    anomalies = extractAnomalyIndices(self.anomalyLikelihood)

    # Highlight weekends in main chart
    self.highlightChart(weekends, self._mainGraph)

    # Highlight anomalies in anomaly chart
    self.highlightChart(anomalies, self._anomalyGraph)

    maxValue = max(self.allValues)
    self._mainGraph.relim()
    self._mainGraph.axes.set_ylim(0, maxValue + (maxValue * 0.02))

    self._mainGraph.relim()
    self._mainGraph.autoscale_view(True, scaley=False)
    self._anomalyGraph.relim()
    self._anomalyGraph.autoscale_view(True, True, True)

    plt.draw()



  def close(self):
    plt.ioff()
    plt.show()



NuPICOutput.register(NuPICFileOutput)
NuPICOutput.register(NuPICPlotOutput)