From d8043b6098bb9464b054c05649cba505874e0e17 Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Fri, 1 Oct 2021 13:19:29 -0500
Subject: [PATCH] Support png files in plot_inline and Plot.to_ipython_image

---
 openmc/executor.py | 20 +++++--------------
 openmc/plots.py    | 50 ++++++++++++++++++++++++++++++----------------
 2 files changed, 38 insertions(+), 32 deletions(-)

diff --git a/openmc/executor.py b/openmc/executor.py
index e8f0de2c7e..882764730a 100644
--- a/openmc/executor.py
+++ b/openmc/executor.py
@@ -3,6 +3,7 @@ from numbers import Integral
 import subprocess
 
 import openmc
+from .plots import _get_plot_image
 
 
 def _run(args, output, cwd):
@@ -62,10 +63,6 @@ def plot_geometry(output=True, openmc_exec='openmc', cwd='.'):
 def plot_inline(plots, openmc_exec='openmc', cwd='.', convert_exec='convert'):
     """Display plots inline in a Jupyter notebook.
 
-    This function requires that you have a program installed to convert PPM
-    files to PNG files. Typically, that would be `ImageMagick
-    <https://www.imagemagick.org>`_ which includes a `convert` command.
-
     Parameters
     ----------
     plots : Iterable of openmc.Plot
@@ -75,7 +72,8 @@ def plot_inline(plots, openmc_exec='openmc', cwd='.', convert_exec='convert'):
     cwd : str, optional
         Path to working directory to run in
     convert_exec : str, optional
-        Command that can convert PPM files into PNG files
+        Command that can convert PPM files into PNG files. Only used if your
+        OpenMC installation was not compiled against libpng.
 
     Raises
     ------
@@ -83,7 +81,7 @@ def plot_inline(plots, openmc_exec='openmc', cwd='.', convert_exec='convert'):
         If the `openmc` executable returns a non-zero status
 
     """
-    from IPython.display import Image, display
+    from IPython.display import display
 
     if not isinstance(plots, Iterable):
         plots = [plots]
@@ -94,16 +92,8 @@ def plot_inline(plots, openmc_exec='openmc', cwd='.', convert_exec='convert'):
     # Run OpenMC in geometry plotting mode
     plot_geometry(False, openmc_exec, cwd)
 
-    images = []
     if plots is not None:
-        for p in plots:
-            if p.filename is not None:
-                ppm_file = f'{p.filename}.ppm'
-            else:
-                ppm_file = f'plot_{p.id}.ppm'
-            png_file = ppm_file.replace('.ppm', '.png')
-            subprocess.check_call([convert_exec, ppm_file, png_file])
-            images.append(Image(png_file))
+        images = [_get_plot_image(p, convert_exec) for p in plots]
         display(*images)
 
 
diff --git a/openmc/plots.py b/openmc/plots.py
index 0abff20d7b..bf8be6dfc4 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -1,6 +1,7 @@
 from collections.abc import Iterable, Mapping
 from numbers import Real, Integral
 from pathlib import Path
+import shutil
 import subprocess
 from xml.etree import ElementTree as ET
 
@@ -165,6 +166,31 @@ _SVG_COLORS = {
 }
 
 
+def _get_plot_image(plot, convert_exec):
+    from IPython.display import Image
+
+    # Check for .png first
+    stem = plot.filename if plot.filename is not None else f'plot_{plot.id}'
+    png_file = f'{stem}.png'
+    if Path(png_file).exists():
+        return Image(png_file)
+
+    # If .png doesn't exist, try finding .ppm
+    ppm_file = f'{stem}.ppm'
+    if not Path(ppm_file).exists():
+        raise FileNotFoundError(f"Could not find image for plot {plot.id}")
+
+    # Check that 'convert' command exists
+    if shutil.which(convert_exec) is None:
+        raise RuntimeError(
+            f"ImageMagick is needed to convert {ppm_file} to .png format. Please "
+            "install ImageMagick and try again.")
+
+    # Convert .ppm to .png and return image
+    subprocess.check_call([convert_exec, ppm_file, png_file])
+    return Image(png_file)
+
+
 class Plot(IDManagerMixin):
     """Definition of a finite region of space to be plotted.
 
@@ -675,11 +701,9 @@ class Plot(IDManagerMixin):
                          convert_exec='convert'):
         """Render plot as an image
 
-        This method runs OpenMC in plotting mode to produce a bitmap image which
-        is then converted to a .png file and loaded in as an
-        :class:`IPython.display.Image` object. As such, it requires that your
-        model geometry, materials, and settings have already been exported to
-        XML.
+        This method runs OpenMC in plotting mode to produce a .png file. If PNG
+        supported is not enabled, try converting the fallback .ppm file to .png
+        using ImageMagick's convert command.
 
         Parameters
         ----------
@@ -688,7 +712,8 @@ class Plot(IDManagerMixin):
         cwd : str, optional
             Path to working directory to run in
         convert_exec : str, optional
-            Command that can convert PPM files into PNG files
+            Command that can convert PPM files into PNG files. Only used if your
+            OpenMC installation was not compiled against libpng.
 
         Returns
         -------
@@ -696,23 +721,14 @@ class Plot(IDManagerMixin):
             Image generated
 
         """
-        from IPython.display import Image
-
         # Create plots.xml
         Plots([self]).export_to_xml()
 
         # Run OpenMC in geometry plotting mode
         openmc.plot_geometry(False, openmc_exec, cwd)
 
-        # Convert to .png
-        if self.filename is not None:
-            ppm_file = f'{self.filename}.ppm'
-        else:
-            ppm_file = f'plot_{self.id}.ppm'
-        png_file = ppm_file.replace('.ppm', '.png')
-        subprocess.check_call([convert_exec, ppm_file, png_file])
-
-        return Image(png_file)
+        # Return produced image
+        return _get_plot_image(self, convert_exec)
 
 
 class Plots(cv.CheckedList):