diff --git a/src/openmc_source_plotter/core.py b/src/openmc_source_plotter/core.py index 2a056f4..6392b26 100644 --- a/src/openmc_source_plotter/core.py +++ b/src/openmc_source_plotter/core.py @@ -98,7 +98,7 @@ def plot_source_energy( xaxis_units: The units to use for the x axis. Options are 'eV' or 'MeV'. """ - if xaxis_units not in ['eV', 'MeV']: + if xaxis_units not in ["eV", "MeV"]: raise ValueError(f"xaxis_units must be either 'eV' or 'MeV' not {xaxis_units}") if figure is None: @@ -121,8 +121,8 @@ def plot_source_energy( if isinstance(self, openmc.SourceBase): probability = probability * self.strength energy = bin_edges[:-1] - if xaxis_units == 'MeV': - energy = energy /1e6 + if xaxis_units == "MeV": + energy = energy / 1e6 # Plot source energy histogram figure.add_trace( plotly.graph_objects.Scatter( diff --git a/tests/test_core_with_source.py b/tests/test_core_with_source.py index 1b9e0cc..ed54bf0 100644 --- a/tests/test_core_with_source.py +++ b/tests/test_core_with_source.py @@ -47,9 +47,14 @@ def test_energy_plot(test_source): assert isinstance(plot, go.Figure) assert len(plot.data[0]["x"]) == 1 + def test_energy_plot_axis(test_source): - plot = test_source.plot_source_energy(n_samples=10, xaxis_type='log', yaxis_type='linear', xaxis_units='eV') - plot = test_source.plot_source_energy(n_samples=10, xaxis_type='linear', yaxis_type='log', xaxis_units='MeV') + plot = test_source.plot_source_energy( + n_samples=10, xaxis_type="log", yaxis_type="linear", xaxis_units="eV" + ) + plot = test_source.plot_source_energy( + n_samples=10, xaxis_type="linear", yaxis_type="log", xaxis_units="MeV" + ) assert isinstance(plot, go.Figure) assert len(plot.data[0]["x"]) == 1