Merge pull request #405 from nicoloabrate/master

Update sensitivity.py

Author: Andrew Johnson

serpentTools/parsers/sensitivity.py

@@ -5,7 +5,7 @@
 from itertools import product
 
 from numpy import transpose, hstack
-from matplotlib.pyplot import gca
+from matplotlib.pyplot import gca, axvline
 
 from serpentTools.utils.plot import magicPlotDocDecorator, formatPlot
 from serpentTools.engines import KeywordParser
@@ -304,8 +304,8 @@ def _postcheck(self):
                                         "stored on reader")
 
     @magicPlotDocDecorator
-    def plot(self, resp, zai=None, pert=None, mat=None, sigma=3,
-             normalize=True, ax=None, labelFmt=None,
+    def plot(self, resp, zai=None, pert=None, mat=None, mevscale=False,
+             egrid=None, sigma=3, normalize=True, ax=None, labelFmt=None,
              title=None, logx=True, logy=False, loglog=False, xlabel=None,
              ylabel=None, legend=None, ncol=1):
         """
@@ -331,6 +331,12 @@ def plot(self, resp, zai=None, pert=None, mat=None, sigma=3,
         mat: None or str or list of strings
             Plot sensitivities due to these materials. Passing ``None``
             will plot against all materials.
+        mevscale : bool, optional
+            Flag for plotting energy grid in MeV units. If ``True``, the energy
+            axis is expressed in MeV. Default is ``False``.
+        egrid : numpy.array, optional
+            User-defined energy grid boundaries displayed on the sensitivities
+            as vblack, dashed vertical lines. Default is ``None``.
         {sigma}
         normalize: True
             Normalize plotted data per unit lethargy
@@ -389,7 +395,7 @@ def plot(self, resp, zai=None, pert=None, mat=None, sigma=3,
 
         errors = resMat[..., 1] * values * sigma
 
-        energies = self.energies * 1E6
+        energies = self.energies if mevscale else self.energies * 1E6
         for z, m, p in product(zais, mats, perts):
             iZ = self.zais[z]
             iM = self.materials[m]
@@ -402,6 +408,13 @@ def plot(self, resp, zai=None, pert=None, mat=None, sigma=3,
             ax.errorbar(energies, yVals, yErrs, label=label,
                         drawstyle='steps-post')
 
+        if egrid is not None:
+            for group in egrid:
+                ax.axvline(group, color='k', linestyle='dashed')
+
+        if xlabel is None:
+            xlabel = "Energy [MeV]" if mevscale else "Energy [eV]"
+
         xlabel = 'Energy [eV]' if xlabel is None else xlabel
         ylabel = ylabel if ylabel is not None else (
             'Sensitivity {} {}'.format(