Support png files in plot_inline and Plot.to_ipython_image

This commit is contained in:
Paul Romano 2021-10-01 13:19:29 -05:00
parent 6eea13043f
commit d8043b6098
2 changed files with 38 additions and 32 deletions

View file

@ -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)

View file

@ -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):