Repository: CORE-GATECH-GROUP/serpent-tools
Pull request #311 (2019-05-15 17:36:06)

TST/DEV Add utility for testing plot functions
In order to increase our coverage, and have some better testing of plots outside of the jupyter notebook (which are really just do they or don't they work), I've added a little test infrastructure. The primary function is `s/t/utils.testPlotAttrs` which is capable of inspection various attributes related to a plot object. Specifically an `axes` subplot instance, as these are often returned from our plot methods. Axis labels and scales, legend labels, and plot titles are candidates for testing. Added some tests specifically responsible for testing this function. 

Demonstrated it's use bu added some tests for the DepletedMaterial and ResultsReader plot methods. Can expand this further, but had some time while waiting on other results to arrive.


Patch:
diff --git a/serpentTools/parsers/results.py b/serpentTools/parsers/results.py
index ebcbcca..e0ecfd6 100644
--- a/serpentTools/parsers/results.py
+++ b/serpentTools/parsers/results.py
@@ -30,6 +30,7 @@ from serpentTools.utils import (
     placeLegend,
     magicPlotDocDecorator,
     deconvertVariableName,
+    RESULTS_PLOT_XLABELS,
 )
 from serpentTools.messages import (
     warning, SerpentToolsException,
@@ -586,8 +587,6 @@ class ResultsReader(XSReader):
         if y is None:
             y = x
             x = "burnDays"
-        if x == 'days':
-            x = 'burnDays'
 
         sigma = max(int(sigma), 0)
 
@@ -597,11 +596,7 @@ class ResultsReader(XSReader):
             right = self._expandPlotIterables(right, ' [right]')
 
         if xlabel is True:
-            xlabel = {
-                'burnup': 'Burnup [MWd/kgU]',
-                'burnDays': 'Burnup [d]',
-                'burnStep': 'Burnup step',
-            }[x]
+            xlabel = RESULTS_PLOT_XLABELS[x]
 
         if len(y) == 1 and ylabel is None:
             for ylabel in y.values():
diff --git a/serpentTools/utils/plot.py b/serpentTools/utils/plot.py
index 0132087..e4cfb54 100644
--- a/serpentTools/utils/plot.py
+++ b/serpentTools/utils/plot.py
@@ -11,6 +11,7 @@ from serpentTools.utils.docstrings import magicPlotDocDecorator
 __all__ = [
     'DETECTOR_PLOT_LABELS',
     'DEPLETION_PLOT_LABELS',
+    'RESULTS_PLOT_XLABELS',
     'formatPlot',
     'addColorbar',
     'setAx_xlims',
@@ -64,6 +65,11 @@ for key in _DET_LABEL_INDEXES:
 
 del _DET_LABEL_INDEXES
 
+RESULTS_PLOT_XLABELS = {
+    'burnup': DEPLETION_PLOT_LABELS['burnup'],
+    'burnDays': DEPLETION_PLOT_LABELS['days'],
+    'burnStep': "Burnup step",
+}
 
 PLOT_FORMAT_DEFAULTS = {
     'xlabel': None, 'ylabel': None, 'legend': True,


Test patch:
diff --git a/serpentTools/tests/test_ResultsReader.py b/serpentTools/tests/test_ResultsReader.py
index 001408b..9f4e678 100644
--- a/serpentTools/tests/test_ResultsReader.py
+++ b/serpentTools/tests/test_ResultsReader.py
@@ -12,6 +12,11 @@ from serpentTools.settings import rc
 from serpentTools.data import getFile, readDataFile
 from serpentTools.parsers import ResultsReader
 from serpentTools.messages import SerpentToolsException
+from serpentTools.utils import RESULTS_PLOT_XLABELS
+from serpentTools.tests.utils import (
+    plotTest,
+    testPlotAttrs,
+)
 
 
 GCU_START_STR = "GCU_UNIVERSE_NAME"
@@ -640,6 +645,60 @@ class RestrictedResultsReader(Serp2129Helper):
 del TesterCommonResultsReader
 
 
+class ResPlotTester(TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        with rc:
+            rc['xs.variableExtras'] = [
+                'ABS_KEFF',
+                'TOT_CPU_TIME',
+                'BURN_DAYS',
+                'BURNUP',
+                'BURN_STEP',
+            ]
+            cls.reader = ResultsReader(getFile('InnerAssembly_res.m'))
+            cls.reader.read()
+
+    @plotTest
+    def test_singlePlot(self):
+        """Test the plot capabilities of the ResultsReader"""
+        ax = self.reader.plot('absKeff', sigma=3)
+        testPlotAttrs(
+            self, ax, xlabel=RESULTS_PLOT_XLABELS['burnDays'],
+            ylabel="absKeff$ \\pm 3\\sigma$",
+        )
+
+        ax.clear()
+        newLabel = 'Multiplication factor'
+        self.reader.plot('burnup', {'absKeff': newLabel}, ax=ax, sigma=0,
+                         logx=True)
+        testPlotAttrs(
+            self, ax, ylabel=newLabel, xscale='log',
+            xlabel=RESULTS_PLOT_XLABELS['burnup'],
+        )
+        # plot two quantities
+        ax.clear()
+        self.reader.plot('burnStep', ['absKeff', 'totCpuTime'], ax=ax,
+                         sigma=0, ylabel="ylabel")
+        testPlotAttrs(
+            self, ax, xlabel=RESULTS_PLOT_XLABELS['burnStep'],
+            ylabel="ylabel", legendLabels=['absKeff', 'totCpuTime'],
+        )
+
+    @plotTest
+    def test_rightPlot(self):
+        """Test plotting on left and right y axis"""
+        left, right = self.reader.plot('absKeff', sigma=0, right='totCpuTime')
+        testPlotAttrs(
+            self, left, ylabel='absKeff',
+            legendLabels=['absKeff', 'totCpuTime [right]'],
+        )
+        testPlotAttrs(
+            self, right, ylabel='totCpuTime',
+        )
+
+
 class ResADFTester(TestCase):
 
     @classmethod
diff --git a/serpentTools/tests/test_depletion.py b/serpentTools/tests/test_depletion.py
index bdcc57a..40c4b43 100644
--- a/serpentTools/tests/test_depletion.py
+++ b/serpentTools/tests/test_depletion.py
@@ -13,7 +13,11 @@ from serpentTools.parsers.depletion import (
     DepletionReader, getMaterialNameAndVariable, getMatlabVarName,
     prepToMatlab, deconvert,
 )
-from serpentTools.tests.utils import LoggerMixin, MatlabTesterHelper
+from serpentTools.utils import DEPLETION_PLOT_LABELS
+from serpentTools.tests.utils import (
+    LoggerMixin, MatlabTesterHelper,
+    plotTest, testPlotAttrs,
+)
 
 
 DEP_FILE = 'ref_dep.m'
@@ -113,6 +117,34 @@ class DepletionTester(_DepletionTestHelper):
             fromGetItem = self.reader[name]
             self.assertIs(mat, fromGetItem, msg=mat)
 
+    @plotTest
+    def test_plotFewIso(self):
+        """Test the basic functionality of the depletion plot"""
+        mat = self.reader[self.MATERIAL]
+        ax = mat.plot('days', 'adens', names='U235')
+        testPlotAttrs(
+            self, ax, xlabel=DEPLETION_PLOT_LABELS['days'],
+            ylabel=DEPLETION_PLOT_LABELS['adens'],
+            xscale='linear', yscale='linear',
+            legendLabels=[],
+        )
+        # clear the plot for a second go
+
+        ax.clear()
+        mat.plot('burnup', 'adens', names=['U235', 'Xe135'], loglog=True)
+        testPlotAttrs(
+            self, ax, xlabel=DEPLETION_PLOT_LABELS['burnup'],
+            xscale='log', yscale='log', legendLabels=['U235', 'Xe135'],
+        )
+
+    @plotTest
+    def test_plotFormatting(self):
+        mat = self.reader[self.MATERIAL]
+        ax = mat.plot('days', 'adens', names='U235', legend=True,
+                      labelFmt="{mat}-{iso}")
+        testPlotAttrs(
+            self, ax, legendLabels=[self.MATERIAL + '-U235'])
+
 
 class DepletedMaterialTester(_DepletionTestHelper):
     """Class that tests the functionality of the DepletedMaterial class"""
diff --git a/serpentTools/tests/test_utils.py b/serpentTools/tests/test_utils.py
index ebb8e6e..6d4b8af 100644
--- a/serpentTools/tests/test_utils.py
+++ b/serpentTools/tests/test_utils.py
@@ -6,7 +6,7 @@ from unittest import TestCase
 
 from numpy import arange, ndarray, array, ones, ones_like, zeros_like
 from numpy.testing import assert_array_equal
-from six import iteritems
+from six import iteritems, assertRaisesRegex
 
 from serpentTools.utils import (
     convertVariableName,
@@ -24,8 +24,11 @@ from serpentTools.utils import (
     DC_STAT_NOT_IMPLEMENTED,
     DC_STAT_DIFF_TYPES,
     DC_STAT_DIFF_SHAPES,
+    formatPlot,
 )
 
+from serpentTools.tests.utils import plotTest, testPlotAttrs
+
 
 class VariableConverterTester(TestCase):
     """Class for testing our variable name conversion function."""
@@ -387,6 +390,58 @@ class SplitDictionaryTester(TestCase):
         self.assertSetEqual(self.goodKeys, goodKeys)
 
 
+class PlotTestTester(TestCase):
+    """Class to test the validity of the plot test utils"""
+    XLABEL = "Test x label"
+    YLABEL = "Test y label"
+
+    def buildPlot(self):
+        from matplotlib.pyplot import gca
+        ax = gca()
+        ax.plot([1, 2, 3], label=1)
+        ax.legend()
+        ax.set_ylabel(self.YLABEL)
+        ax.set_xlabel(self.XLABEL)
+        return ax
+
+    @plotTest
+    def test_plotAttrs_gca(self):
+        """Test the testPlotAttrs without passing an axes object"""
+        self.buildPlot()
+        testPlotAttrs(self, xlabel=self.XLABEL, ylabel=self.YLABEL)
+
+    @plotTest
+    def test_plotAttrs_fail(self):
+        """Test the failure modes of the testPlotAttrs function"""
+        ax = self.buildPlot()
+        with assertRaisesRegex(self, AssertionError, 'xlabel'):
+            testPlotAttrs(self, ax, xlabel="Bad label")
+        with assertRaisesRegex(self, AssertionError, 'ylabel'):
+            testPlotAttrs(self, ax, ylabel="Bad label")
+        with assertRaisesRegex(self, AssertionError, 'xscale'):
+            testPlotAttrs(self, ax, xscale="log")
+        with assertRaisesRegex(self, AssertionError, 'yscale'):
+            testPlotAttrs(self, ax, yscale="log")
+        with assertRaisesRegex(self, AssertionError, 'legend text'):
+            testPlotAttrs(self, ax, legendLabels='bad text')
+        with assertRaisesRegex(self, AssertionError, 'legend text'):
+            testPlotAttrs(self, ax, legendLabels=['1', '2'])
+        with assertRaisesRegex(self, AssertionError, 'title'):
+            testPlotAttrs(self, ax, title='Bad title')
+
+    @plotTest
+    def test_formatPlot(self):
+        """Test the capabilities of the formatPlot function"""
+        ax = self.buildPlot()
+        newX = 'new x label'
+        newY = 'new y label'
+        title = 'plot title'
+        formatPlot(ax, xlabel=newX, ylabel=newY, loglog=True,
+                   title=title)
+        testPlotAttrs(self, ax, xlabel=newX, ylabel=newY,
+                      xscale='log', yscale='log', title=title)
+
+
 if __name__ == '__main__':
     from unittest import main
     main()
diff --git a/serpentTools/tests/utils.py b/serpentTools/tests/utils.py
index c8fcbc3..c654589 100644
--- a/serpentTools/tests/utils.py
+++ b/serpentTools/tests/utils.py
@@ -193,6 +193,44 @@ def plotTest(f):
 
 
 def getLegendTexts(ax):
-    """Return all texts for items in legend."""
-    lgd = ax.legend()
+    """
+    Return all texts for items in legend.
+
+    An empty list signifies no legend or no labeled
+    objects.
+    """
+    lgd = ax.get_legend()
+
+    if lgd is None:
+        return []
+
     return [item.get_text() for item in lgd.get_texts()]
+
+
+def testPlotAttrs(testobj, ax=None, xlabel=None, ylabel=None,
+                  xscale=None, yscale=None, legendLabels=None,
+                  title=None):
+    """Compare properties of a generated axes object"""
+    if ax is None:
+        from matplotlib.pyplot import gca
+        ax = gca()
+
+    if xlabel is not None:
+        testobj.assertEqual(xlabel, ax.get_xlabel(), msg='xlabel')
+
+    if ylabel is not None:
+        testobj.assertEqual(ylabel, ax.get_ylabel(), msg='ylabel')
+
+    if xscale is not None:
+        testobj.assertEqual(xscale, ax.get_xscale(), msg='xscale')
+
+    if yscale is not None:
+        testobj.assertEqual(yscale, ax.get_yscale(), msg='yscale')
+
+    if legendLabels is not None:
+        if isinstance(legendLabels, str):
+            legendLabels = [legendLabels, ]
+        testobj.assertEqual(legendLabels, getLegendTexts(ax),
+                            msg='legend text')
+    if title is not None:
+        testobj.assertEqual(title, ax.get_title(), msg='title')