Simple historian for SImpy simulations

SimpyHistorian.py

class Historian():
  """A simple historian class for logging simulation data.
  
  Args: 
    item_id (str): a name for the identifier of the tracked objects
    data_callbacks (dict): a dictionary of name:function pairs where name 
      is a string label for the data element being recorded, and the function
      returns the current value for the data element.
  """
  
  def __init__(self, item_id, data_callbacks = {'time':lambda: env.now}):
    self.item_id = item_id
    self.data_callbacks = data_callbacks
    self.log = []
  
  def write_log(self, item_id, event):
    data_record = [item_id, event]
    data_record.extend([fcn() for fcn in self.data_callbacks.values()])
    self.log.append(data_record)
    
  def data(self, name):
    columns = [self.item_id, 'event']
    columns.extend(self.data_callbacks.keys())
    df = pd.DataFrame(self.log, columns=columns).set_index(self.item_id)
    return df.pivot(columns='event', values=name)
  
  def plot(self, *names):
    names = names if names else self.data_callbacks.keys()
    N = len(names)
    n = 1
    for name in names:
      ax = plt.subplot(N, 1, n)
      self.data(name).plot(ax=ax, title=name)
      plt.legend(loc='center left', bbox_to_anchor=(1.0, 0.5))
      n += 1
    plt.tight_layout()