API Support spectrumPlot(..., by=None|"max") (#457)

* STY Apply some black formatting to spectrumPlot

* BUG Use logy in Detector.spectrumPlot

* DEV Use broadcasting to improve spectrum normalization

Avoids for loops and lets numpy handle division at once

* API Support spectrumPlot(..., by={None, "max"})

Default is "max" which will scale tallied data by the max value
in each column (plotted data). Passing by=None will not perform
any additional modifications. Neither options changes the division
by lethargy.

Closes #455

* TST Test spectrumPlot(..., by=None)

Author: drewejohnson

serpentTools/detectors.py

@@ -27,7 +27,7 @@
 
 from numpy import (
     unique, inf, hstack, arange, log, divide, asfortranarray,
-    ndarray, asarray,
+    ndarray, asarray, newaxis
 )
 from matplotlib.patches import RegularPolygon
 from matplotlib.collections import PatchCollection
@@ -495,10 +495,25 @@ def _compare(self, other, lower, upper, sigma):
         return similar
 
     @magicPlotDocDecorator
-    def spectrumPlot(self, fixed=None, ax=None, normalize=True, xlabel=None,
-                     ylabel=None, steps=True, logx=True, logy=False,
-                     loglog=False, sigma=3, labels=None, legend=None, ncol=1,
-                     title=None, **kwargs):
+    def spectrumPlot(
+            self,
+            fixed=None,
+            ax=None,
+            normalize=True,
+            by="max",
+            xlabel=None,
+            ylabel=None,
+            steps=True,
+            logx=True,
+            logy=False,
+            loglog=False,
+            sigma=3,
+            labels=None,
+            legend=None,
+            ncol=1,
+            title=None,
+            **kwargs,
+    ):
         """
         Quick plot of the detector value as a function of energy.
 
@@ -508,7 +523,12 @@ def spectrumPlot(self, fixed=None, ax=None, normalize=True, xlabel=None,
             Dictionary controlling the reduction in data
         {ax}
         normalize: bool
-            Normalize quantities per unit lethargy
+            Divide tally values by size of lethargy bins.
+        by : {"max", None}, optional
+            How to scale the spectrum after considering ``normalize``.
+            ``"max"`` (default) will divide by the maximum value such
+            that the max value is unity. ``None`` will not perform any
+            additional modifications.
         {xlabel}
         {ylabel}
         steps: bool
@@ -546,13 +566,17 @@ def spectrumPlot(self, fixed=None, ax=None, normalize=True, xlabel=None,
             )
         elif len(slicedTallies.shape) == 1:
             slicedTallies = slicedTallies.reshape(slicedTallies.size, 1)
+
         lowerE = self.grids['E'][:, 0]
         if normalize:
             lethBins = log(
                 divide(self.grids['E'][:, -1], lowerE))
-            for indx in range(slicedTallies.shape[1]):
-                scratch = divide(slicedTallies[:, indx], lethBins)
-                slicedTallies[:, indx] = scratch / scratch.max()
+            # Need to scale up the dimension of the energy grids to support
+            # broadcasting -> (E, N) / (E, )
+            slicedTallies /= lethBins[:, newaxis]
+
+        if by == "max":
+            slicedTallies /= slicedTallies.max(axis=0)
 
         if steps:
             if 'drawstyle' in kwargs:
@@ -576,9 +600,17 @@ def spectrumPlot(self, fixed=None, ax=None, normalize=True, xlabel=None,
 
         if legend is None and labels:
             legend = True
-        ax = formatPlot(ax, loglog=loglog, logx=logx, legendcols=ncol,
-                        xlabel=xlabel or "Energy [MeV]", ylabel=ylabel,
-                        legend=legend, title=title)
+        ax = formatPlot(
+            ax,
+            loglog=loglog,
+            logx=logx,
+            logy=logy,
+            legendcols=ncol,
+            xlabel=xlabel or "Energy [MeV]",
+            ylabel=ylabel,
+            legend=legend,
+            title=title,
+        )
         return ax
 
     @magicPlotDocDecorator

tests/plots/result_images/testSpectrumJustNormPerLethargy.png

@@ -17,3 +17,9 @@ def testBwrMeshPlot():
     with pytest.warns(None) as record:
         reader["xymesh"].meshPlot(fixed={"energy": 0})
     assert len(record) == 0
+
+
+@compare_or_update_plot
+def testSpectrumJustNormPerLethargy():
+    reader = readDataFile("bwr_det0.m")
+    reader["spectrum"].spectrumPlot(by=None, fixed={"reaction": 1})