From 4cadcbcf029310ced0eb3e7f638259f600335238 Mon Sep 17 00:00:00 2001 From: Jonathan Shimwell Date: Fri, 28 Apr 2023 04:26:21 +0100 Subject: [PATCH] more flexible cross section plotting (#2478) --------- Co-authored-by: Paul Romano --- openmc/plotter.py | 193 +++++++++++++++++-------------- tests/unit_tests/test_plotter.py | 83 ++++++++++++- 2 files changed, 189 insertions(+), 87 deletions(-) diff --git a/openmc/plotter.py b/openmc/plotter.py index 0362a8d176..d8e06e4fe1 100644 --- a/openmc/plotter.py +++ b/openmc/plotter.py @@ -1,6 +1,5 @@ from itertools import chain from numbers import Integral, Real -import string import numpy as np @@ -56,7 +55,49 @@ _MAX_E = 20.e6 ELEMENT_NAMES = list(openmc.data.ELEMENT_SYMBOL.values())[1:] -def plot_xs(this, types, divisor_types=None, temperature=294., axis=None, +def _get_legend_label(this, type): + """Gets a label for the element or nuclide or material and reaction plotted""" + if isinstance(this, str): + return f'{this} {type}' + elif this.name is '': + return f'Material {this.id} {type}' + else: + return f'{this.name} {type}' + + +def _get_yaxis_label(reactions, divisor_types): + """Gets a y axis label for the type of data plotted""" + + if all(isinstance(item, str) for item in reactions.keys()): + stem = 'Microscopic' + if divisor_types: + mid, units = 'Data', '' + else: + mid, units = 'Cross Section', '[b]' + elif all(isinstance(item, openmc.Material) for item in reactions.keys()): + stem = 'Macroscopic' + if divisor_types: + mid, units = 'Data', '' + else: + mid, units = 'Cross Section', '[1/cm]' + else: + msg = "Mixture of openmc.Material and elements/nuclides. Invalid type for plotting" + raise TypeError(msg) + + return f'{stem} {mid} {units}' + + +def _get_title(reactions): + """Gets a title for the type of data plotted""" + if len(reactions) == 1: + this, = reactions + name = this.name if isinstance(this, openmc.Material) else this + return f'Cross Section Plot For {name}' + else: + return 'Cross Section Plot' + + +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): @@ -64,10 +105,10 @@ def plot_xs(this, types, divisor_types=None, temperature=294., axis=None, Parameters ---------- - this : str or openmc.Material - Object to source data from. Nuclides and elements can be input as a str - types : Iterable of values of PLOT_TYPES - The type of cross sections to include in the plot. + reactions : dict + keys can be either a nuclide or element in string form or an + openmc.Material object. Values are the type of cross sections to + include in the plot. divisor_types : Iterable of values of PLOT_TYPES, optional Cross section types which will divide those produced by types before plotting. A type of 'unity' can be used to effectively not @@ -115,49 +156,6 @@ def plot_xs(this, types, divisor_types=None, temperature=294., axis=None, import matplotlib.pyplot as plt cv.check_type("plot_CE", plot_CE, bool) - cv.check_type("this", this, (str, openmc.Material)) - - if plot_CE: - # Calculate for the CE cross sections - E, data = calculate_cexs(this, types, temperature, sab_name, - ce_cross_sections, enrichment) - if divisor_types: - cv.check_length('divisor types', divisor_types, len(types)) - Ediv, data_div = calculate_cexs(this, divisor_types, temperature, - sab_name, ce_cross_sections, - enrichment) - - # Create a new union grid, interpolate data and data_div on to that - # grid, and then do the actual division - Enum = E[:] - E = np.union1d(Enum, Ediv) - data_new = np.zeros((len(types), len(E))) - - for line in range(len(types)): - data_new[line, :] = \ - np.divide(np.interp(E, Enum, data[line, :]), - np.interp(E, Ediv, data_div[line, :])) - if divisor_types[line] != 'unity': - types[line] = types[line] + ' / ' + divisor_types[line] - data = data_new - else: - # Calculate for MG cross sections - E, data = calculate_mgxs(this, types, orders, temperature, - mg_cross_sections, ce_cross_sections, - enrichment) - - if divisor_types: - cv.check_length('divisor types', divisor_types, len(types)) - Ediv, data_div = calculate_mgxs(this, divisor_types, - divisor_orders, temperature, - mg_cross_sections, - ce_cross_sections, enrichment) - - # Perform the division - for line in range(len(types)): - data[line, :] /= data_div[line, :] - if divisor_types[line] != 'unity': - types[line] += ' / ' + divisor_types[line] # Generate the plot if axis is None: @@ -165,18 +163,69 @@ def plot_xs(this, types, divisor_types=None, temperature=294., axis=None, else: fig = None ax = axis + + all_types = [] + + for this, types in reactions.items(): + all_types = all_types + types + + if plot_CE: + cv.check_type("this", this, (str, openmc.Material)) + # Calculate for the CE cross sections + E, data = calculate_cexs(this, types, temperature, sab_name, + ce_cross_sections, enrichment) + if divisor_types: + cv.check_length('divisor types', divisor_types, len(types)) + Ediv, data_div = calculate_cexs(this, divisor_types, temperature, + sab_name, ce_cross_sections, + enrichment) + + # Create a new union grid, interpolate data and data_div on to that + # grid, and then do the actual division + Enum = E[:] + E = np.union1d(Enum, Ediv) + data_new = np.zeros((len(types), len(E))) + + for line in range(len(types)): + data_new[line, :] = \ + np.divide(np.interp(E, Enum, data[line, :]), + np.interp(E, Ediv, data_div[line, :])) + if divisor_types[line] != 'unity': + types[line] = types[line] + ' / ' + divisor_types[line] + data = data_new + else: + # Calculate for MG cross sections + E, data = calculate_mgxs(this, types, orders, temperature, + mg_cross_sections, ce_cross_sections, + enrichment) + + if divisor_types: + cv.check_length('divisor types', divisor_types, len(types)) + Ediv, data_div = calculate_mgxs(this, divisor_types, + divisor_orders, temperature, + mg_cross_sections, + ce_cross_sections, enrichment) + + # Perform the division + for line in range(len(types)): + data[line, :] /= data_div[line, :] + if divisor_types[line] != 'unity': + types[line] += ' / ' + divisor_types[line] + + # Plot the data + for i in range(len(data)): + data[i, :] = np.nan_to_num(data[i, :]) + if np.sum(data[i, :]) > 0.: + ax.plot(E, data[i, :], label=_get_legend_label(this, types[i])) + # Set to loglog or semilogx depending on if we are plotting a data # type which we expect to vary linearly - if set(types).issubset(PLOT_TYPES_LINEAR): - plot_func = ax.semilogx + if set(all_types).issubset(PLOT_TYPES_LINEAR): + ax.set_xscale('log') + ax.set_yscale('linear') else: - plot_func = ax.loglog - - # Plot the data - for i in range(len(data)): - data[i, :] = np.nan_to_num(data[i, :]) - if np.sum(data[i, :]) > 0.: - plot_func(E, data[i, :], label=types[i]) + ax.set_xscale('log') + ax.set_yscale('log') ax.set_xlabel('Energy [eV]') if plot_CE: @@ -184,33 +233,9 @@ def plot_xs(this, types, divisor_types=None, temperature=294., axis=None, else: ax.set_xlim(E[-1], E[0]) - if divisor_types: - if isinstance(this, str): - if this in ELEMENT_NAMES: - ylabel = 'Elemental Microscopic Data' - else: - ylabel = 'Nuclide Microscopic Data' - elif isinstance(this, openmc.Material): - ylabel = 'Macroscopic Data' - else: - raise TypeError("Invalid type for plotting") - else: - if isinstance(this, str): - if this in ELEMENT_NAMES: - ylabel = 'Elemental Cross Section [b]' - else: - ylabel = 'Microscopic Cross Section [b]' - elif isinstance(this, openmc.Material): - ylabel = 'Macroscopic Cross Section [1/cm]' - else: - raise TypeError("Invalid type for plotting") - ax.set_ylabel(ylabel) + ax.set_ylabel(_get_yaxis_label(reactions, divisor_types)) ax.legend(loc='best') - name = this.name if isinstance(this, openmc.Material) else this - if len(types) > 1: - ax.set_title('Cross Sections for ' + name) - else: - ax.set_title('Cross Section for ' + name) + ax.set_title(_get_title(reactions)) return fig diff --git a/tests/unit_tests/test_plotter.py b/tests/unit_tests/test_plotter.py index 8520249ba8..8a1147b4af 100644 --- a/tests/unit_tests/test_plotter.py +++ b/tests/unit_tests/test_plotter.py @@ -11,6 +11,7 @@ def test_mat(): mat_1.add_element("C", 4.0, "ao") return mat_1 + def test_calculate_cexs_elem_mat_sab(test_mat): """Checks that sab cross sections are included in the _calculate_cexs_elem_mat method and have the correct shape""" @@ -74,9 +75,85 @@ def test_calculate_cexs_with_materials(test_mat): @pytest.mark.parametrize("this", ["Be", "Be9"]) def test_plot_xs(this): from matplotlib.figure import Figure - assert isinstance(openmc.plot_xs(this, types=['total']), 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, types=['total']), Figure) + assert isinstance(openmc.plot_xs({test_mat: ['total']}), Figure) + +def test_plot_axes_labels(): + # just nuclides + axis_label = openmc.plotter._get_yaxis_label( + reactions={ + 'Li6': [205], + 'Li7': [205], + }, divisor_types=False + ) + assert axis_label == 'Microscopic Cross Section [b]' + + # just elements + axis_label = openmc.plotter._get_yaxis_label( + reactions={ + 'Li': [205], + 'Be': [16], + }, divisor_types=False + ) + assert axis_label == 'Microscopic Cross Section [b]' + + # mixed nuclide and element + axis_label = openmc.plotter._get_yaxis_label( + reactions={ + 'Li': [205], + 'Li7': [205], + }, divisor_types=False + ) + assert axis_label == 'Microscopic Cross Section [b]' + + # just materials + mat1 = openmc.Material() + mat1.add_nuclide('Fe56', 1) + mat1.set_density('g/cm3', 1) + mat2 = openmc.Material() + mat2.add_element('Fe', 1) + mat2.add_nuclide('Fe55', 1) + mat2.set_density('g/cm3', 1) + axis_label = openmc.plotter._get_yaxis_label( + reactions={ + mat1: [205], + mat2: [16], + }, divisor_types=False + ) + assert axis_label == 'Macroscopic Cross Section [1/cm]' + + # mixed materials and nuclides + with pytest.raises(TypeError): + openmc.plotter._get_yaxis_label( + reactions={'Li6': [205], mat2: [16]}, + divisor_types=False + ) + + # mixed materials and elements + with pytest.raises(TypeError): + openmc.plotter._get_yaxis_label( + reactions={'Li': [205], mat2: [16]}, + divisor_types=False + ) + + +def test_get_title(): + title = openmc.plotter._get_title(reactions={'Li': [205]}) + assert title == 'Cross Section Plot For Li' + title = openmc.plotter._get_title(reactions={'Li6': [205]}) + assert title == 'Cross Section Plot For Li6' + title = openmc.plotter._get_title(reactions={ + 'Li6': [205], + 'Li7': [205] + }) + assert title == 'Cross Section Plot' + + mat1 = openmc.Material() + mat1.add_nuclide('Fe56', 1) + 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' \ No newline at end of file