Adding energy axis units to plot xs (openmc-dev#2876)

Co-authored-by: Paul Romano <[email protected]>

openmc/plotter.py

@@ -105,7 +105,7 @@ def _get_title(reactions):
 def plot_xs(reactions, divisor_types=None, temperature=294., axis=None,
             sab_name=None, ce_cross_sections=None, mg_cross_sections=None,
             enrichment=None, plot_CE=True, orders=None, divisor_orders=None,
-            **kwargs):
+            energy_axis_units="eV", **kwargs):
     """Creates a figure of continuous-energy cross sections for this item.
 
     Parameters
@@ -148,6 +148,10 @@ def plot_xs(reactions, divisor_types=None, temperature=294., axis=None,
     **kwargs :
         All keyword arguments are passed to
         :func:`matplotlib.pyplot.figure`.
+    energy_axis_units : {'eV', 'keV', 'MeV'}
+        Units used on the plot energy axis
+
+        .. versionadded:: 0.14.1
 
     Returns
     -------
@@ -161,6 +165,9 @@ def plot_xs(reactions, divisor_types=None, temperature=294., axis=None,
     import matplotlib.pyplot as plt
 
     cv.check_type("plot_CE", plot_CE, bool)
+    cv.check_value("energy_axis_units", energy_axis_units, {"eV", "keV", "MeV"})
+
+    axis_scaling_factor = {"eV": 1.0, "keV": 1e-3, "MeV": 1e-6}
 
     # Generate the plot
     if axis is None:
@@ -217,6 +224,8 @@ def plot_xs(reactions, divisor_types=None, temperature=294., axis=None,
                     if divisor_types[line] != 'unity':
                         types[line] += ' / ' + divisor_types[line]
 
+        E *= axis_scaling_factor[energy_axis_units]
+
         # Plot the data
         for i in range(len(data)):
             data[i, :] = np.nan_to_num(data[i, :])
@@ -232,9 +241,12 @@ def plot_xs(reactions, divisor_types=None, temperature=294., axis=None,
         ax.set_xscale('log')
         ax.set_yscale('log')
 
-    ax.set_xlabel('Energy [eV]')
+    ax.set_xlabel(f"Energy [{energy_axis_units}]")
     if plot_CE:
-        ax.set_xlim(_MIN_E, _MAX_E)
+        ax.set_xlim(
+            _MIN_E * axis_scaling_factor[energy_axis_units],
+            _MAX_E * axis_scaling_factor[energy_axis_units],
+        )
     else:
         ax.set_xlim(E[-1], E[0])
 

tests/unit_tests/test_plotter.py

@@ -77,10 +77,19 @@ def test_plot_xs(this):
     from matplotlib.figure import Figure
     assert isinstance(openmc.plot_xs({this: ['total', 'elastic']}), Figure)
 
+
 def test_plot_xs_mat(test_mat):
     from matplotlib.figure import Figure
     assert isinstance(openmc.plot_xs({test_mat: ['total']}), Figure)
 
+
+@pytest.mark.parametrize("units", ["eV", "keV", "MeV"])
+def test_plot_xs_energy_axis(units):
+    plot = openmc.plot_xs({'Be9': ['(n,2n)']}, energy_axis_units=units)
+    axis_text = plot.get_axes()[0].get_xaxis().get_label().get_text()
+    assert axis_text == f'Energy [{units}]'
+
+
 def test_plot_axes_labels():
     # just nuclides
     axis_label = openmc.plotter._get_yaxis_label(
@@ -156,4 +165,4 @@ def test_get_title():
     mat1.set_density('g/cm3', 1)
     mat1.name = 'my_mat'
     title = openmc.plotter._get_title(reactions={mat1: [205]})
-    assert title == 'Cross Section Plot For my_mat'
+    assert title == 'Cross Section Plot For my_mat'