Repository: mroberge/hydrofunctions

Issue:
Add ability to modify the legend and titles
* HydroFunctions version: 0.1.8dev
* Python version: 3.7

### Description

Hydrofunction's built-in charts should have the ability to set a title or create a legend.


### Potential Solutions:

1) add new parameters `legend` and `title`

    - `legend` default could be `False`; otherwise you could provide a value for legend.
    - This approach might require a `legend_loc` parameter too.
    - `title` could be set to `False` or text.

2) Use **kwargs and pass these on
3) What does Pandas do??

Comments:

Patch:
diff --git a/hydrofunctions/charts.py b/hydrofunctions/charts.py
index 8e41052..aba1b8e 100644
--- a/hydrofunctions/charts.py
+++ b/hydrofunctions/charts.py
@@ -20,7 +20,7 @@ from matplotlib.ticker import NullFormatter
 import numpy as np
 
 
-def flow_duration(Qdf, xscale='logit', yscale='log', ylabel='Stream Discharge (m³/s)', symbol='.', legend=True):
+def flow_duration(Qdf, xscale='logit', yscale='log', ylabel='Stream Discharge (m³/s)', symbol='.', legend=True, legend_loc='best', title=''):
     """Creates a flow duration chart from a dataframe of discharges.
 
     Args:
@@ -56,6 +56,17 @@ def flow_duration(Qdf, xscale='logit', yscale='log', ylabel='Stream Discharge (m
 
         legend (bool, default: True): whether the legend should be plotted.
 
+        legend_loc (str, default: 'best'): the location of the legend.
+
+            * 'best': Automatically choose the option below with the least overlap.
+            * 'upper left', 'upper right', 'lower left', 'lower right': place the legend at the corresponding corner of the axes/figure.
+            * 'upper center', 'lower center', 'center left', 'center right': place the legend at the center of the corresponding edge of the axes/figure.
+            * 'center': place the legend at the center of the axes/figure.
+            * The location can also be a 2-tuple giving the coordinates of the lower-left corner of the legend in axes coordinates.
+
+        title (str, default: ''): text to use as a figure title. If no text
+            is provided, no title will be created (default).
+
     Returns:
         fig (matplotlib.figure.Figure):
             a matplotlib figure. This will plot immediately in a Jupyter
@@ -76,14 +87,16 @@ def flow_duration(Qdf, xscale='logit', yscale='log', ylabel='Stream Discharge (m
     ax.set_yscale(yscale)
     ax.set_ylabel(ylabel)
     if legend:
-        ax.legend()
+        ax.legend(loc=legend_loc)
+    if title:
+        ax.title.set_text(title)
     # A pyplot bug causes a valueError value if the xlabel is set.
     #ax.set_xlabel('Probability of Exceedence')
     ax.xaxis.set_minor_formatter(NullFormatter())
     return fig, ax
 
 
-def cycleplot(Qseries, cycle='diurnal', compare=None, y_label='Discharge (ft³/s)'):
+def cycleplot(Qseries, cycle='diurnal', compare=None, y_label='Discharge (ft³/s)', legend=True, legend_loc='best', title=''):
     """Creates a chart to illustrate annual and diurnal cycles.
 
     This chart will use the pandas groupby method to plot the mean and median
@@ -138,6 +151,19 @@ def cycleplot(Qseries, cycle='diurnal', compare=None, y_label='Discharge (ft³/s
 
         y_label (str): The label for the y axis.
 
+        legend (bool, default: True): whether the legend should be plotted.
+
+        legend_loc (str, default: 'best'): the location of the legend.
+
+            * 'best': Automatically choose the option below with the least overlap.
+            * 'upper left', 'upper right', 'lower left', 'lower right': place the legend at the corresponding corner of the axes/figure.
+            * 'upper center', 'lower center', 'center left', 'center right': place the legend at the center of the corresponding edge of the axes/figure.
+            * 'center': place the legend at the center of the axes/figure.
+            * The location can also be a 2-tuple giving the coordinates of the lower-left corner of the legend in axes coordinates.
+
+        title (str, default: ''): text to use as a figure title. If no text
+            is provided, no title will be created (default).
+
     Returns:
         fig (matplotlib.figure.Figure):
             a matplotlib figure. This will plot immediately in a Jupyter
@@ -187,9 +213,7 @@ def cycleplot(Qseries, cycle='diurnal', compare=None, y_label='Discharge (ft³/s
         cycleby = Qseries.index.hour
         x_label = ' (hour of the day)'
     else:
-        print("The cycle label '", cycle, "' is not recognized as an option. Using cycle='diurnal' instead.")
-        cycleby = Qseries.index.hour
-        x_label = ' (hour of the day)'
+        raise ValueError("The cycle label '", cycle, "' is not recognized as an option.")
 
     if compare is None:
         # Don't make a comparison plot.
@@ -269,7 +293,8 @@ def cycleplot(Qseries, cycle='diurnal', compare=None, y_label='Discharge (ft³/s
         # axs[i].xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%H'))
 
     # Set the legend on either the ax or fig.
-    axs[0].legend(loc='best', fancybox=True, framealpha=0.5)
+    if legend:
+        axs[0].legend(loc=legend_loc, fancybox=True, framealpha=0.5)
     # fig.legend(loc='upper center', shadow=True, frameon=True, fancybox=True, framealpha=0.5)
 
     # Get the yaxis limits, set bottom to zero.
@@ -278,5 +303,7 @@ def cycleplot(Qseries, cycle='diurnal', compare=None, y_label='Discharge (ft³/s
     axs[0].set_ylabel(y_label)
     axs[0].set_xlabel('Time' + x_label)
     plt.tight_layout()
+    if title:
+        fig.suptitle(title)
 
     return fig, axs


Test patch:
diff --git a/tests/test_charts.py b/tests/test_charts.py
index 9219578..a74174d 100644
--- a/tests/test_charts.py
+++ b/tests/test_charts.py
@@ -34,18 +34,28 @@ class TestFlowDuration(unittest.TestCase):
         actual_yscale = actual_ax.yaxis.get_scale()
         actual_ylabel = actual_ax.yaxis.get_label_text()
         actual_marker = actual_ax.get_lines()[0].get_marker()
+        actual_legend = actual_ax.get_legend()
+        actual_legend_loc = actual_legend._loc
+        actual_title = actual_ax.get_title()
 
         self.assertEqual(actual_xscale, 'logit')
         self.assertEqual(actual_yscale, 'log')
         self.assertEqual(actual_ylabel, 'Stream Discharge (m³/s)')
         self.assertEqual(actual_marker, '.')
+        self.assertTrue(actual_legend)
+        self.assertEqual(actual_legend_loc, 0) # '0' is internal code for 'best'.
+        self.assertEqual(actual_title, '')
 
     def test_charts_flowduration_accepts_params(self):
         expected = pd.DataFrame(data=dummy)
         params = {'xscale': 'linear',
                   'yscale': 'linear',
                   'ylabel': 'test value',
-                  'symbol': ','}
+                  'symbol': ',',
+                  'legend': False,
+                  'legend_loc': 'center',
+                  'title': 'Test Title',
+                  }
 
         actual_fig, actual_ax = charts.flow_duration(expected, **params)
 
@@ -53,12 +63,19 @@ class TestFlowDuration(unittest.TestCase):
         actual_yscale = actual_ax.yaxis.get_scale()
         actual_ylabel = actual_ax.yaxis.get_label_text()
         actual_marker = actual_ax.get_lines()[0].get_marker()
+        actual_legend = actual_ax.get_legend()
+        # There is no legend in this test, so there is no legend property.
+        #actual_legend_loc = actual_legend._loc
+        actual_title = actual_ax.get_title()
 
         self.assertEqual(actual_xscale, 'linear')
         self.assertEqual(actual_yscale, 'linear')
         self.assertEqual(actual_ylabel, 'test value')
         self.assertEqual(actual_marker, ',')
-
+        self.assertIsNone(actual_legend)
+        # There is no legend, so there is no legend location property.
+        #self.assertEqual(actual_legend_loc, 10) # 'center' is equal to 10.
+        self.assertEqual(actual_title, 'Test Title')
 
 class TestCyclePlot(unittest.TestCase):
 
@@ -80,15 +97,21 @@ class TestCyclePlot(unittest.TestCase):
     def test_charts_cycleplot_parts(self):
         expected_df, expected_dict = hf.extract_nwis_df(test_json, interpolate=False)
 
-        actual_fig, actual_ax = charts.cycleplot(expected_df)
+        actual_fig, actual_ax = charts.cycleplot(expected_df, legend_loc='center', title='test title')
 
         actual_xscale = actual_ax[0].xaxis.get_scale()
         actual_yscale = actual_ax[0].yaxis.get_scale()
         actual_ylabel = actual_ax[0].yaxis.get_label_text()
+        actual_legend = actual_ax[0].get_legend()
+        actual_legend_loc = actual_legend._loc
+        actual_title = actual_fig._suptitle.get_text() # unofficial title accessor! A little wonky.
 
         self.assertEqual(actual_xscale, 'linear')
         self.assertEqual(actual_yscale, 'linear')
         self.assertEqual(actual_ylabel, 'Discharge (ft³/s)')
+        self.assertTrue(actual_legend)
+        self.assertEqual(actual_legend_loc, 10) # '10' is internal code for legend(loc = 'center')
+        self.assertEqual(actual_title, 'test title')
 
     def test_charts_cycleplot_compare_month(self):
         expected_df, expected_dict = hf.extract_nwis_df(test_json, interpolate=False)
@@ -102,5 +125,46 @@ class TestCyclePlot(unittest.TestCase):
         self.assertIsInstance(actual_fig, matplotlib.figure.Figure)
         self.assertIsInstance(actual_ax[0], matplotlib.axes.Axes)
 
+    def test_charts_cycleplot_cycle_annual(self):
+        expected_df, expected_dict = hf.extract_nwis_df(test_json, interpolate=False)
+        actual_fig, actual_ax = charts.cycleplot(expected_df, 'annual')
+        self.assertIsInstance(actual_fig, matplotlib.figure.Figure)
+        self.assertIsInstance(actual_ax[0], matplotlib.axes.Axes)
+
+    def test_charts_cycleplot_cycle_annualdate(self):
+        expected_df, expected_dict = hf.extract_nwis_df(test_json, interpolate=False)
+        actual_fig, actual_ax = charts.cycleplot(expected_df, 'annual-date')
+        self.assertIsInstance(actual_fig, matplotlib.figure.Figure)
+        self.assertIsInstance(actual_ax[0], matplotlib.axes.Axes)
+
+    def test_charts_cycleplot_cycle_annualmonth(self):
+        expected_df, expected_dict = hf.extract_nwis_df(test_json, interpolate=False)
+        actual_fig, actual_ax = charts.cycleplot(expected_df, 'annual-month')
+        self.assertIsInstance(actual_fig, matplotlib.figure.Figure)
+        self.assertIsInstance(actual_ax[0], matplotlib.axes.Axes)
+
+    def test_charts_cycleplot_cycle_weekly(self):
+        expected_df, expected_dict = hf.extract_nwis_df(test_json, interpolate=False)
+        actual_fig, actual_ax = charts.cycleplot(expected_df, 'weekly')
+        self.assertIsInstance(actual_fig, matplotlib.figure.Figure)
+        self.assertIsInstance(actual_ax[0], matplotlib.axes.Axes)
+
+    def test_charts_cycleplot_cycle_diurnalsmallest(self):
+        expected_df, expected_dict = hf.extract_nwis_df(test_json, interpolate=False)
+        actual_fig, actual_ax = charts.cycleplot(expected_df, 'diurnal-smallest')
+        self.assertIsInstance(actual_fig, matplotlib.figure.Figure)
+        self.assertIsInstance(actual_ax[0], matplotlib.axes.Axes)
+
+    def test_charts_cycleplot_cycle_diurnalhour(self):
+        expected_df, expected_dict = hf.extract_nwis_df(test_json, interpolate=False)
+        actual_fig, actual_ax = charts.cycleplot(expected_df, 'diurnal-hour')
+        self.assertIsInstance(actual_fig, matplotlib.figure.Figure)
+        self.assertIsInstance(actual_ax[0], matplotlib.axes.Axes)
+
+    def test_charts_cycleplot_cycle_nonsense_raises_ValueError(self):
+        expected_df, expected_dict = hf.extract_nwis_df(test_json, interpolate=False)
+        with self.assertRaises(ValueError):
+            actual_fig, actual_ax = charts.cycleplot(expected_df, 'nonsense')
+
 if __name__ == '__main__':
     unittest.main(verbosity=2)