diff --git a/openmc/universe.py b/openmc/universe.py index b8abf663d6..7195688292 100644 --- a/openmc/universe.py +++ b/openmc/universe.py @@ -1,3 +1,4 @@ +import math import typing from abc import ABC, abstractmethod from collections import OrderedDict @@ -300,7 +301,7 @@ class Universe(UniverseBase): _default_legend_kwargs = {'bbox_to_anchor': ( 1.05, 1), 'loc': 2, 'borderaxespad': 0.0} - def plot(self, origin=None, width=None, pixels=(200, 200), + def plot(self, origin=None, width=None, pixels=40000, basis='xy', color_by='cell', colors=None, seed=None, openmc_exec='openmc', axes=None, legend=False, legend_kwargs=_default_legend_kwargs, outline=False, @@ -319,8 +320,12 @@ class Universe(UniverseBase): universe.bounding_box.width() will be used to attempt to ascertain the plot width. Defaults to (10, 10) if the bounding_box contains inf values - pixels : Iterable of int - Number of pixels to use in each basis direction + pixels : Iterable of int or int + If iterable of ints provided then this directly sets the number of + pixels to use in each basis direction. If int provided then this + sets the total number of pixels in the plot and the number of + pixels in each basis direction is calculated from this total and + the image aspect ratio. basis : {'xy', 'xz', 'yz'} The basis directions for the plot color_by : {'cell', 'material'} @@ -398,6 +403,11 @@ class Universe(UniverseBase): y_width = bb_width['xyz'.index(basis[1])] width = (x_width, y_width) + if isinstance(pixels, int): + aspect_ratio = width[0] / width[1] + pixels_y = math.sqrt(pixels / aspect_ratio) + pixels = (int(pixels / pixels_y), int(pixels_y)) + x_min = origin[x] - 0.5*width[0] x_max = origin[x] + 0.5*width[0] y_min = origin[y] - 0.5*width[1] diff --git a/tests/unit_tests/test_universe.py b/tests/unit_tests/test_universe.py index aedc814383..f4cf43d919 100644 --- a/tests/unit_tests/test_universe.py +++ b/tests/unit_tests/test_universe.py @@ -81,7 +81,7 @@ def test_plot(run_in_tmpdir, sphere_model): colors=colors, color_by="cell", legend=False, - pixels=(10, 10), + pixels=100, basis=basis, outline=False )