Repository: watts-dev/watts (language: python)

Issue:
Support for OpenMC geometry plotting / volume calculations
Right now, the `OpenMCDriver` class only calls `openmc.run` but it would nice to also have the ability to call other OpenMC functions, such as `openmc.plot_geometry` and `openmc.calculate_volumes` during invocation of the `OpenMCDriver`.


Discussion / hints:

Fix patch:
diff --git a/doc/source/user/usage.rst b/doc/source/user/usage.rst
index 41bc66b8..cc15ba5c 100644
--- a/doc/source/user/usage.rst
+++ b/doc/source/user/usage.rst
@@ -194,6 +194,20 @@ the plugin are forwarded to the :func:`openmc.run` function. For example::
 
     results = openmc_plugin(params, mpi_args=["mpiexec", "-n", "16"])
 
+By default, the OpenMC plugin will only call the :func:`openmc.run` function,
+but you can customize the execution by passing an arbitrary function as the
+``function`` keyword argument. For example, if you wanted to additionally call
+:func:`openmc.plot_geometry` each time the plugin is called, this could be
+accomplished as follows::
+
+    import openmc
+
+    def run_function():
+        openmc.plot_geometry()
+        openmc.run()
+
+    results = openmc_plugin(params, function=run_function)
+
 PyARC Plugin
 ~~~~~~~~~~~~~
 
diff --git a/examples/1App_OpenMC_VHTR/watts_exec.py b/examples/1App_OpenMC_VHTR/watts_exec.py
index c99d23a2..19ae1bae 100644
--- a/examples/1App_OpenMC_VHTR/watts_exec.py
+++ b/examples/1App_OpenMC_VHTR/watts_exec.py
@@ -2,11 +2,11 @@
 # SPDX-License-Identifier: MIT
 
 """
-This example demonstrates how to use WATTS to perform 
+This example demonstrates how to use WATTS to perform
 OpenMC calculation. This example uses a simple VHTR unit-cell
 model with 1 coolant channel surrounded by graphite and fuel.
 The demonstration includes the application of unit-conversion
-approach in WATTS. OpenMC is executed and the main results are 
+approach in WATTS. OpenMC is executed and the main results are
 printed out and stored in the params database.
 """
 
@@ -15,6 +15,7 @@
 import watts
 from statistics import mean
 from openmc_template import build_openmc_model
+import openmc
 from astropy.units import Quantity
 
 
@@ -53,14 +54,23 @@
 
 params.show_summary(show_metadata=False, sort_by='time')
 
-# Run OpenMC plugin
+# Create OpenMC plugin
 openmc_plugin = watts.PluginOpenMC(build_openmc_model, show_stderr=True) # show only error
-openmc_result = openmc_plugin(params)
+
+# Run OpenMC plugin, instructing it to plot the geometry and run a simulation
+def run_func():
+    openmc.plot_geometry()
+    openmc.run()
+openmc_result = openmc_plugin(params, function=run_func)
 print("KEFF = ", openmc_result.keff)
 print(openmc_result.inputs)
 print(openmc_result.outputs)
 print(openmc_result.tallies[0].get_pandas_dataframe())
 
+# Open folder in order to view plot images that were produced
+openmc_result.open_folder()
+
+# Set power fractions in parameters
 power_fractions = openmc_result.tallies[0].get_values(scores=['nu-fission']).ravel()
 for i, power_frac in enumerate(power_fractions):
     params[f'Init_P_{i+1}'] = power_frac
diff --git a/src/watts/plugin_openmc.py b/src/watts/plugin_openmc.py
index ad69473d..5f3036da 100644
--- a/src/watts/plugin_openmc.py
+++ b/src/watts/plugin_openmc.py
@@ -57,7 +57,10 @@ def keff(self):
         # Get k-effective from last statepoint
         last_statepoint = self.statepoints[-1]
         with openmc.StatePoint(last_statepoint) as sp:
-            return sp.k_combined
+            if hasattr(sp, 'keff'):
+                return sp.keff
+            else:
+                return sp.k_combined
 
     @property
     @lru_cache()
@@ -113,13 +116,21 @@ def prerun(self, params: Parameters) -> None:
         if self.model_builder is not None:
             self.model_builder(params_copy)
 
-    def run(self, **kwargs: Mapping):
+    def run(self, function: Optional[Callable] = None, **kwargs: Mapping):
         """Run OpenMC
 
         Parameters
         ----------
+        function
+            Function to execute. If not passed, defaults to only calling
+            :func:`openmc.run`.
         **kwargs
-            Keyword arguments passed on to :func:`openmc.run`
+            Keyword arguments passed on to ``function``
+
+        See also
+        --------
+        openmc.run, openmc.plot_geometry, openmc.calculate_volumes
+
         """
         print("Run for OpenMC Plugin")
         import openmc
@@ -127,7 +138,10 @@ def run(self, **kwargs: Mapping):
             func_stdout = tee_stdout if self.show_stdout else redirect_stdout
             func_stderr = tee_stderr if self.show_stderr else redirect_stderr
             with func_stdout(f), func_stderr(f):
-                openmc.run(**kwargs)
+                if function:
+                    function(**kwargs)
+                else:
+                    openmc.run(**kwargs)
 
     def postrun(self, params: Parameters, name: str) -> ResultsOpenMC:
         """Collect information from OpenMC simulation and create results object
@@ -168,8 +182,11 @@ def files_since(pattern, time):
         outputs = ['OpenMC_log.txt']
         outputs.extend(files_since('tallies.out', self._run_time))
         outputs.extend(files_since('source.*.h5', self._run_time))
-        outputs.extend(files_since('particle.*.h5', self._run_time))
+        outputs.extend(files_since('particle*.h5', self._run_time))
         outputs.extend(files_since('statepoint.*.h5', self._run_time))
+        outputs.extend(files_since('volume*.h5', self._run_time))
+        outputs.extend(files_since('*.png', self._run_time))
+        outputs.extend(files_since('*.ppm', self._run_time))
 
         time = datetime.fromtimestamp(self._run_time * 1e-9)
         return ResultsOpenMC(params, name, time, inputs, outputs)


Test patch:
diff --git a/tests/test_plugin_openmc.py b/tests/test_plugin_openmc.py
index 968fc02a..54dc123d 100644
--- a/tests/test_plugin_openmc.py
+++ b/tests/test_plugin_openmc.py
@@ -1,6 +1,9 @@
 # SPDX-FileCopyrightText: 2022 UChicago Argonne, LLC
 # SPDX-License-Identifier: MIT
 
+from pathlib import Path
+
+import pytest
 import openmc
 import watts
 
@@ -73,8 +76,9 @@ def test_openmc_plugin():
     assert last_result.keff.s == result.keff.s
 
 
-def test_extra_inputs(run_in_tmpdir):
-    # Create inputs manually -- 10 cm sphere of pure Pu239
+@pytest.fixture
+def pu_model():
+    """Model of a 10 cm sphere of Pu239"""
     model = openmc.model.Model()
     mat = openmc.Material()
     mat.add_nuclide('Pu239', 1.0)
@@ -85,7 +89,12 @@ def test_extra_inputs(run_in_tmpdir):
     model.settings.batches = 10
     model.settings.inactive = 0
     model.settings.particles = 1000
-    model.export_to_xml()
+    return model
+
+
+def test_extra_inputs(run_in_tmpdir, pu_model):
+    # Export model to XML
+    pu_model.export_to_xml()
 
     # Use OpenMC plugin with extra inputs
     plugin = watts.PluginOpenMC(extra_inputs=['geometry.xml', 'materials.xml', 'settings.xml'])
@@ -94,3 +103,22 @@ def test_extra_inputs(run_in_tmpdir):
 
     input_names = {p.name for p in result.inputs}
     assert input_names == {'geometry.xml', 'materials.xml', 'settings.xml'}
+
+
+def test_arbitrary_function(run_in_tmpdir, pu_model):
+    # Add a plot to the model and export
+    plot = openmc.Plot()
+    plot.filename = 'watts_plot'
+    plot.pixels = (100, 100)
+    plot.width = (20.0, 20.0)
+    pu_model.plots.append(plot)
+    pu_model.export_to_xml()
+
+    # Use OpenMC plugin to run a geometry plot
+    plugin = watts.PluginOpenMC(extra_inputs=Path.cwd().glob('*.xml'))
+    params = watts.Parameters()
+    result = plugin(params, function=openmc.plot_geometry)
+
+    # Generated plot should be in outputs
+    output_names = {p.name for p in result.outputs}
+    assert 'watts_plot.png' in output_names