From 5e47706d088bfe75de30777cf92082d1d4a0e106 Mon Sep 17 00:00:00 2001
From: Matt Hoffman <mwh30@cam.ac.uk>
Date: Wed, 3 Jun 2015 01:08:57 +0100
Subject: [PATCH] Update plotting.

---
 mwhutils/plotting.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/mwhutils/plotting.py b/mwhutils/plotting.py
index 17c36b6..258a8bb 100644
--- a/mwhutils/plotting.py
+++ b/mwhutils/plotting.py
@@ -92,14 +92,17 @@ def scatter(self, x, y, marker='o', color='k', alpha=1):
         self._ax.scatter(x, y, **kwargs)
         self._draw()
 
-    def plot(self, x, y):
+    def plot(self, x, y=None, color=None, alpha=1):
         """
         Add a simple plot to the axis.
         """
-        self._ax.plot(x, y)
+        if y is None:
+            y = x
+            x = np.arange(len(y))
+        self._ax.plot(x, y, lw=2, color=color, alpha=alpha)
         self._draw()
 
-    def plot_banded(self, x, y, a=None, b=None):
+    def plot_banded(self, x, y, a=None, b=None, label=''):
         if a is None and b is None:
             lo = np.zeros_like(y)
             hi = y
@@ -109,7 +112,7 @@ def plot_banded(self, x, y, a=None, b=None):
         else:
             lo = y - a
             hi = y + b
-        lines = self._ax.plot(x, y)
+        lines = self._ax.plot(x, y, lw=2, label=label)
         color = lines[0].get_color()
         alpha = 0.25
         self._ax.fill_between(x, lo, hi, color=color, alpha=alpha)
@@ -132,6 +135,7 @@ def hline(self, y):
     def _draw(self):
         self._ax.axis('tight')
         self._ax.axis(self._lim)
+        self._ax.legend()
         self._hook()
         if not self._hold and pl.isinteractive():
             self.draw()