Repository: CORE-GATECH-GROUP/serpent-tools
Pull request #229 (2018-08-13 15:48:13)

[API/TST] Better x/y labels for sensitivity plots; Test sensitivity plot
Previously there was not way to not set the x/y labels for a sensitivity plot. If the passed value did not evaluate to True, then the default was used. This mainly occurred during subplotting routines. Now, passing an argument that does not evaluate to True will not trigger the default x/y label, and result in not setting the label.

Also added some testing for the sensitivity plotter. Mainly for label formatting and x/y axis label setting.

- [x] PR fits the [project scope](http://serpent-tools.readthedocs.io/en/latest/develop/contributing.html#project-scope)
- [x] Code to be committed is [DRY](https://en.wikipedia.org/wiki/Don%27t_repeat_yourself) and matches the [code style](http://serpent-tools.readthedocs.io/en/latest/develop/codestyle.html#code-style) of the project
- [x] New and/or changed features are [well documented](http://serpent-tools.readthedocs.io/en/latest/develop/documentation.html#documentation) and have adequate testing

## TODO
- [x] Add this PR to changelog
- [x] Remove the useless `titleFmt` argument for the plot


Patch:
diff --git a/docs/changelog.rst b/docs/changelog.rst
index dba3651..507d3d6 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -15,6 +15,9 @@ Next
 * :pull:`227` - Better documentation of our :ref:`cli`.
   Better documentation and testing of functions for generating input
   files with unique random seeds - :mod:`serpentTools.seed`
+* :pull:`229` - :meth:`serpentTools.parsers.sensitivity.SensitivityReader.plot`
+  now respects the option to not set x nor y labels.
+  
 
 .. _v0.5.2:
 
diff --git a/serpentTools/parsers/sensitivity.py b/serpentTools/parsers/sensitivity.py
index b14f628..c668efa 100644
--- a/serpentTools/parsers/sensitivity.py
+++ b/serpentTools/parsers/sensitivity.py
@@ -235,7 +235,7 @@ class SensitivityReader(BaseReader):
 
     @magicPlotDocDecorator
     def plot(self, resp, zai=None, pert=None, mat=None, sigma=3,
-             normalize=True, ax=None, labelFmt=None, titleFmt=None,
+             normalize=True, ax=None, labelFmt=None,
              title=None, logx=True, logy=False, loglog=False, xlabel=None,
              ylabel=None, legend=True, ncol=1):
         """
@@ -332,23 +332,23 @@ class SensitivityReader(BaseReader):
             ax.errorbar(energies, yVals, yErrs, label=label,
                         drawstyle='steps-post')
 
-        xlabel = xlabel or 'Energy [eV]'
-        ylabel = ylabel or (
+        xlabel = 'Energy [eV]' if xlabel is None else xlabel
+        ylabel = ylabel if ylabel is not None else (
             'Sensitivity {} {}'.format(
                 'per unit lethargy' if normalize else '',
                 r'$\pm{}\sigma$'.format(sigma) if sigma else ''))
         ax = formatPlot(ax, loglog=loglog, logx=logx, logy=logy, ncol=ncol,
-                        legend=legend, xlabel=xlabel, ylabel=ylabel)
+                        legend=legend, xlabel=xlabel, ylabel=ylabel.strip())
         return ax
 
     def _getCleanedPertOpt(self, key, value):
         """Return a set of all or some of the requested perturbations."""
         assert hasattr(self, key), key
-        opts = set(getattr(self, key).keys())
+        opts = getattr(self, key).keys()
         if value is None:
-            return opts
+            return list(opts)
         requested = set([value, ]) if isinstance(value, str) else set(value)
-        missing = {str(xx) for xx in requested.difference(opts)}
+        missing = {str(xx) for xx in requested.difference(set(opts))}
         if missing:
             raise KeyError("Could not find the following perturbations: "
                            "{}".format(', '.join(missing)))


Test patch:
diff --git a/serpentTools/tests/test_sensivity.py b/serpentTools/tests/test_sensivity.py
index dce7b1b..59a4ea7 100644
--- a/serpentTools/tests/test_sensivity.py
+++ b/serpentTools/tests/test_sensivity.py
@@ -1,8 +1,9 @@
 """
 Test the sensitivity reader
 """
-import unittest
+from unittest import TestCase
 from collections import OrderedDict
+from itertools import product
 
 from six import iteritems
 from numpy import array, inf
@@ -11,16 +12,20 @@ from numpy.testing import assert_allclose
 from serpentTools.data import getFile
 from serpentTools.parsers.sensitivity import SensitivityReader
 from serpentTools.tests import compareDictOfArrays
+from serpentTools.tests.utils import (plotTest, getLegendTexts)
 
 TEST_FILE = getFile('bwr_sens0.m')
 
 
-class SensitivityTester(unittest.TestCase):
-    """Class for testing the sensitivity reader."""
+class SensitivityTestHelper(TestCase):
     def setUp(self):
         self.reader = SensitivityReader(TEST_FILE)
         self.reader.read()
 
+
+class SensitivityTester(SensitivityTestHelper):
+    """Class for testing the sensitivity reader."""
+
     def test_expectedSensitivities(self):
         """Verify the sensitivity arrays are loaded correctly."""
         expected = {
@@ -194,5 +199,89 @@ class SensitivityTester(unittest.TestCase):
         self.assertDictEqual(expected, actual)
 
 
+class SensitivityPlotTester(SensitivityTestHelper):
+    """Class for testing rudimentary plot aspects of the reader."""
+
+    ZAIS = [922350]
+    RESP = 'keff'
+    DEFAULT_XLABEL = "Energy [eV]"
+
+    def _plot(self, **kwargs):
+        """Shortcut for plotting."""
+        return self.reader.plot(self.RESP, **kwargs)
+
+    def _checkAxisLabels(self, ax, xlabel, ylabel, msg=None):
+        self.assertEqual(ax.get_xlabel(), xlabel, msg=msg)
+        self.assertEqual(ax.get_ylabel(), ylabel, msg=msg)
+
+    @plotTest
+    def test_plot_normalized(self):
+        """Verify the default axis labels when normalized and with errors."""
+        ax = self._plot(normalize=True, sigma=3)
+        self._checkAxisLabels(ax, self.DEFAULT_XLABEL,
+                              'Sensitivity per unit lethargy $\\pm3\\sigma$')
+
+    @plotTest
+    def test_plot_notNormalized_noSigma(self):
+        """
+        Verify the default axis labels when not normalized nor errors given.
+        """
+        ax = self._plot(normalize=False, sigma=0)
+        self._checkAxisLabels(ax, self.DEFAULT_XLABEL,
+                              'Sensitivity')
+
+    @plotTest
+    def test_plot_passLabels_noset(self):
+        """Verify that the axis labels can be set to empty strings."""
+        xlabel = False
+        ylabel = ""
+        nosetLabel = ""  # un-set matplotlib axis label
+        ax = self._plot(xlabel=xlabel, ylabel=ylabel)
+        self._checkAxisLabels(ax, nosetLabel, nosetLabel)
+
+    def _generateLegendTexts(self, formatter, zais=None, mats=None,
+                             perts=None):
+        reader = self.reader
+        zais = zais or list(reader.zais.keys())
+        mats = mats or list(reader.materials.keys())
+        perts = perts or list(reader.perts.keys())
+        texts = []
+        # match iteration order of sensitivity reader
+        for z, m, p in product(zais, mats, perts):
+            texts.append(formatter.format(z=z, m=m, p=p, r=self.RESP))
+        return texts
+
+    @plotTest
+    def test_plot_labelFormatter(self):
+        """Verify the label formatter for sensitivity plots."""
+        labelFmt = "{m} {z} {p} {r}"
+        ax = self._plot(labelFmt=labelFmt)
+        actual = getLegendTexts(ax)
+        expected = self._generateLegendTexts(labelFmt)
+        self.assertListEqual(actual, expected)
+
+    @plotTest
+    def test_plot_labelFormatter_oneIso(self):
+        """Verify the label formatter for sensitivity plots - pass one ZAI."""
+        labelFmt = "{m} {z} {p} {r}"
+        ax = self._plot(zai=922350, labelFmt=labelFmt)
+        actual = getLegendTexts(ax)
+        expected = self._generateLegendTexts(labelFmt, zais=[922350, ])
+        self.assertListEqual(actual, expected)
+
+    @plotTest
+    def test_plot_raiseError_missingPert(self):
+        """Verify that an error is raised if a bad perturbation is passed."""
+        with self.assertRaises(KeyError):
+            self._plot(zai=-100)
+
+    @plotTest
+    def test_plot_raiseError_missingResp(self):
+        """Verify that an error is raised if a bad response is passed."""
+        with self.assertRaises(KeyError):
+            self.reader.plot("THIS SHOULD FAIL")
+
+
 if __name__ == '__main__':
-    unittest.main()
+    from unittest import main
+    main()
diff --git a/serpentTools/tests/utils.py b/serpentTools/tests/utils.py
index 69ee4d5..5fcc4ae 100644
--- a/serpentTools/tests/utils.py
+++ b/serpentTools/tests/utils.py
@@ -2,6 +2,7 @@
 Utilities to make testing easier
 """
 
+from functools import wraps
 from unittest import TestCase
 from logging import NOTSET
 
@@ -163,3 +164,21 @@ class TestCaseWithLogCapture(TestCase, LoggerMixin):
         failMsg = "Found {} match for {} under {} but should not have"
         self.assertFalse(self.msgInLogs(level, msg, partial),
                          msg=failMsg.format(matchType, msg, level))
+
+
+def plotTest(f):
+    """Decorator that clears up existing plots prior to test."""
+    from matplotlib.pyplot import close, figure
+
+    @wraps(f)
+    def wrappedTest(*args, **kwargs):
+        close('all')
+        figure()
+        return f(*args, **kwargs)
+    return wrappedTest
+
+
+def getLegendTexts(ax):
+    """Return all texts for items in legend."""
+    lgd = ax.legend()
+    return [item.get_text() for item in lgd.get_texts()]