diff --git a/kf_book/book_plots.py b/kf_book/book_plots.py
index 8b61b8c..e6fb685 100644
--- a/kf_book/book_plots.py
+++ b/kf_book/book_plots.py
@@ -528,20 +528,18 @@ def plot_track(xs, ys=None, label='Track', c='k', lw=2, **kwargs):
 
 
 def plot_filter(xs, ys=None, c='#013afe', label='Filter', var=None, **kwargs):
-#def plot_filter(xs, ys=None, c='#6d904f', label='Filter', vars=None, **kwargs):
-
-
+    """ plot result of KF with color `c`, optionally displaying the variance
+    of `xs`. Returns the list of lines generated by plt.plot()"""
+ 
     if ys is None:
         ys = xs
         xs = range(len(ys))
 
-    plt.plot(xs, ys, color=c, label=label, **kwargs)
-
+    lines = plt.plot(xs, ys, color=c, label=label, **kwargs)
     if var is None:
-        return
+        return lines
 
     var = np.asarray(var)
-
     std = np.sqrt(var)
     std_top = ys+std
     std_btm = ys-std
@@ -551,6 +549,8 @@ def plot_filter(xs, ys=None, c='#013afe', label='Filter', var=None, **kwargs):
     plt.fill_between(xs, std_btm, std_top,
                      facecolor='yellow', alpha=0.2)
 
+    return lines
+
 
 
 
@@ -625,4 +625,4 @@ if __name__ == "__main__":
     '''p = [0.2245871, 0.06288015, 0.06109133, 0.0581008, 0.09334062, 0.2245871,
      0.06288015, 0.06109133, 0.0581008,  0.09334062]*2
     bar_plot(p)
-    plot_measurements(p)'''
\ No newline at end of file
+    plot_measurements(p)'''
diff --git a/kf_book/gh_internal.py b/kf_book/gh_internal.py
index b57c6ed..968755f 100644
--- a/kf_book/gh_internal.py
+++ b/kf_book/gh_internal.py
@@ -54,13 +54,13 @@ def plot_gh_results(weights, estimates, predictions, time_step=0):
         plt.gcf().canvas.draw()
         time.sleep(time_step)
 
-        book_plots.plot_filter(xs[:i+1], estimates[:i+1], marker='o')
+        est, = book_plots.plot_filter(xs[:i+1], estimates[:i+1], marker='o')
         plt.xlim([-1, n+1])
         plt.ylim([156.0, 173])
         plt.gcf().canvas.draw()
         time.sleep(time_step)
 
-        plt.legend([act, scale, pred], ['Actual Weight', 'Measurement', 'Predictions'], loc=4)
+        plt.legend([act, scale, est, pred], ['Actual Weight', 'Measurement', 'Estimates', 'Predictions'], loc=4)
     book_plots.set_labels(x='day', y='weight (lbs)')