Auto pixel calc for universe plot (#2513)

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]

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
         )