bluechips/controllers/graph.py

   1 """
   2 Generate various plots
   3 """
   4 import matplotlib
   5 matplotlib.use('Agg')
   6 import matplotlib.pyplot as plt
   7 from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
   8 import numpy as np
   9 import datetime
  10
  11 from cStringIO import StringIO
  12 from pylons import request, response
  13
  14 from bluechips.lib.base import *
  15 import bluechips.lib.totals
  16 from bluechips import model
  17 from bluechips.model import meta, orm, User, Expenditure, Split, Transfer
  18
  19 import time
  20
  21 def get_share(transaction, user):
  22     if type(transaction) == Expenditure:
  23         amount = transaction.share(user)
  24         if transaction.spender == user:
  25             amount -= transaction.amount
  26         return amount
  27     else:
  28         if user == transaction.debtor:
  29             return -transaction.amount
  30         elif user == transaction.creditor:
  31             return transaction.amount
  32         else:
  33             return model.types.Currency(0)
  34
  35 class GraphController(BaseController):
  36     def balance(self):
  37         """
  38         Create a plot showing the net balance for each user over the
  39         course of time.
  40         """
  41         t0 = time.time()
  42         users = meta.Session.query(User).filter(User.resident==True).all()
  43         expenditures = meta.Session.query(Expenditure).\
  44                        options(orm.eagerload_all(Expenditure.splits, Split.user)).\
  45                        all()
  46         transfers = meta.Session.query(Transfer).\
  47                     options(orm.eagerload(Transfer.debtor, Transfer.creditor)).\
  48                     all()
  49
  50         transactions = expenditures + transfers
  51         transactions.sort(key=lambda x: x.date)
  52
  53         t1 = time.time()
  54         totals = dict()
  55         dates = [t.date for t in transactions]
  56         for u in users:
  57             totals[u.name] = np.cumsum([get_share(t,u) for t in transactions])
  58
  59         t2 = time.time()
  60         fig = plt.figure()
  61         canvas = FigureCanvas(fig)
  62         ax = fig.add_subplot(1,1,1)
  63         ax.set_color_cycle(plt.cm.jet(np.linspace(0,1,len(users))))
  64         lines = dict()
  65         ymax = 0
  66         for u,y in totals.items():
  67             lines[u] = plt.step(dates,totals[u], label=u, lw=2)
  68
  69         ax.legend(loc='best')
  70         ax.grid(True)
  71         plt.axis('tight')
  72         ax.set_xlim(xmin=datetime.date.today()-datetime.timedelta(days=365))
  73         ax.set_ylim(ymin=0)
  74         fig.autofmt_xdate(rotation=30)
  75         ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter("%h"))
  76         ax.set_yticklabels(map(lambda x: '%s$%i'%('-' if x<0 else '',abs(x)//100),
  77                                ax.yaxis.get_majorticklocs()))
  78
  79         t3 = time.time()
  80         s = StringIO()
  81         canvas.print_figure(s)
  82         response.headers['Content-Type'] = 'image/png'
  83         t4 = time.time()
  84 #        print t1-t0, t2-t1, t3-t2, t4-t3
  85         return s.getvalue()
  86
  87     def pie(self):
  88         ruser = request.environ['user']
  89
  90         debts = []
  91         debtors = []
  92         credits = []
  93         creditors = []
  94         explode_debt = []
  95         explode_credit = []
  96         eps = model.types.Currency("2.00")
  97         net = bluechips.lib.totals.debts().items()
  98         net.sort(key=lambda x: x[1])
  99         for user, amount in net:
 100             if amount > eps:
 101                 debtors.append(user.name)
 102                 debts.append(amount)
 103                 explode_debt.append(0.1 if user == ruser else 0)
 104             elif amount < -eps:
 105                 creditors.append(user.name)
 106                 credits.append(-amount)
 107                 explode_credit.append(0.1 if user == ruser else 0)
 108
 109         fig = plt.figure(figsize=(6,3))
 110         canvas = FigureCanvas(fig)
 111
 112         ax1 = fig.add_axes([0.05,0.1,0.4,0.8])
 113         ax1.pie(debts,
 114                 explode=explode_debt,
 115                 labels=debtors,
 116                 colors=plt.cm.YlOrRd(np.linspace(0.05,0.95,len(debtors))))
 117         ax1.set_title('Debtors')
 118
 119         ax2 = fig.add_axes([0.55,0.1,0.4,0.8])
 120         ax2.pie(credits,
 121                 explode=explode_credit,
 122                 labels=creditors,
 123                 colors=plt.cm.Blues(np.linspace(0.05,0.95,len(creditors))))
 124         ax2.set_title('Creditors')
 125
 126 #        ax.legend(loc='best')
 127 #        ax.grid(True)
 128 #        plt.axis('tight')
 129 #        ax.set_xlim(xmin=datetime.date.today()-datetime.timedelta(days=365))
 130 #        ax.set_ylim(ymin=0)
 131 #        fig.autofmt_xdate(rotation=30)
 132 #        ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter("%h"))
 133 #        ax.set_yticklabels(map(lambda x: '%s$%i'%('-' if x<0 else '',abs(x)//100),
 134 #                               ax.yaxis.get_majorticklocs()))
 135
 136         s = StringIO()
 137         canvas.print_figure(s, dpi=72)
 138         response.headers['Content-Type'] = 'image/png'
 139         return s.getvalue()