Introduce openmc.lib.TemporarySession context manager (#3475)

Author: paulromano

docs/source/pythonapi/capi.rst

@@ -89,6 +89,7 @@ Classes
    SphericalMesh
    SurfaceFilter
    Tally
+   TemporarySession
    UniverseFilter
    UnstructuredMesh
    WeightFilter

openmc/deplete/chain.py

@@ -11,6 +11,7 @@
 from collections import defaultdict, namedtuple
 from collections.abc import Mapping, Iterable
 from numbers import Real, Integral
+from pathlib import Path
 from warnings import warn
 from typing import List
 
@@ -278,7 +279,6 @@ def __len__(self):
         """Number of nuclides in chain."""
         return len(self.nuclides)
 
-
     @property
     def stable_nuclides(self) -> List[Nuclide]:
         """List of stable nuclides available in the chain"""
@@ -298,6 +298,7 @@ def add_nuclide(self, nuclide: Nuclide):
             Nuclide to add
 
         """
+        _invalidate_chain_cache(self)
         self.nuclide_dict[nuclide.name] = len(self.nuclides)
         self.nuclides.append(nuclide)
 
@@ -463,7 +464,6 @@ def from_endf(cls, decay_files, fpy_files, neutron_files,
                     nuclide.add_reaction('fission', None, q_value, 1.0)
                     fissionable = True
 
-
             if fissionable:
                 if parent in fpy_data:
                     fpy = fpy_data[parent]
@@ -558,6 +558,9 @@ def from_xml(cls, filename, fission_q=None):
             nuc = Nuclide.from_xml(nuclide_elem, root, this_q)
             chain.add_nuclide(nuc)
 
+        # Store path of XML file (used for handling cache invalidation)
+        chain._xml_path = str(Path(filename).resolve())
+
         return chain
 
     def export_to_xml(self, filename):
@@ -888,7 +891,7 @@ def set_branch_ratios(self, branch_ratios, reaction="(n,gamma)",
         --------
         :meth:`get_branch_ratios`
         """
-
+        _invalidate_chain_cache(self)
         # Store some useful information through the validation stage
 
         sums = {}
@@ -1027,6 +1030,7 @@ def fission_yields(self):
 
     @fission_yields.setter
     def fission_yields(self, yields):
+        _invalidate_chain_cache(self)
         if yields is not None:
             if isinstance(yields, Mapping):
                 yields = [yields]
@@ -1249,6 +1253,10 @@ def _follow(self, isotopes, level):
         return found
 
 
+# A global cache for Chain objects
+_CHAIN_CACHE = {}
+
+
 def _get_chain(
     chain_file: PathLike | Chain | None = None,
     fission_q: dict | None = None
@@ -1269,16 +1277,39 @@ def _get_chain(
     Chain
         Depletion chain instance.
     """
+    # If chain_file is already a Chain, return it directly
     if isinstance(chain_file, Chain):
         return chain_file
-    elif isinstance(chain_file, PathLike | None):
-        if chain_file is None:
-            chain_file = openmc.config.get('chain_file')
-            if 'chain_file' not in openmc.config:
-                raise DataError(
-                    "No depletion chain specified and could not find depletion "
-                    "chain in openmc.config['chain_file']"
-                )
-        return Chain.from_xml(chain_file, fission_q)
-    else:
+
+    # Resolve chain_file based on config if None
+    if chain_file is None:
+        chain_file = openmc.config.get('chain_file')
+        if 'chain_file' not in openmc.config:
+            raise DataError(
+                "No depletion chain specified and could not find depletion "
+                "chain in openmc.config['chain_file']"
+            )
+    elif not isinstance(chain_file, PathLike):
         raise TypeError("chain_file must be path-like, a Chain, or None")
+
+    # Determine the key for the cache, which consists of the absolute path, the
+    # file modification time, the file size, and the fission Q values.
+    chain_path = Path(chain_file).resolve()
+    stat_result = chain_path.stat()
+    fq_tuple = tuple(sorted(fission_q.items())) if fission_q else ()
+    key = (chain_path, stat_result.st_mtime, stat_result.st_size, fq_tuple)
+
+    # Check the global cache. If not cached, load the chain from XML and store
+    global _CHAIN_CACHE
+    if key not in _CHAIN_CACHE:
+        _CHAIN_CACHE[key] = Chain.from_xml(chain_path, fission_q)
+    return _CHAIN_CACHE[key]
+
+
+def _invalidate_chain_cache(chain):
+    """Invalidate the cache for a specific Chain (when it is modifed)."""
+    if hasattr(chain, '_xml_path'):
+        # Remove all entries with the same path as self._xml_path
+        for key in list(_CHAIN_CACHE.keys()):
+            if str(key[0]) == chain._xml_path:
+                del _CHAIN_CACHE[key]

openmc/deplete/microxs.py

@@ -14,7 +14,6 @@
 import numpy as np
 
 from openmc.checkvalue import check_type, check_value, check_iterable_type, PathLike
-from openmc.utility_funcs import change_directory
 from openmc import StatePoint
 from openmc.mgxs import GROUP_STRUCTURES
 from openmc.data import REACTION_MT
@@ -286,6 +285,10 @@ def from_multigroup_flux(
         sections available. MicroXS entry will be 0 if the nuclide cross section
         is not found.
 
+        It is recommended to make repeated calls to this method within a context
+        manager using the :class:`openmc.lib.TemporarySession` class to avoid
+        re-initializing OpenMC and loading cross sections each time.
+
         .. versionadded:: 0.15.0
 
         Parameters
@@ -349,37 +352,23 @@ def from_multigroup_flux(
         # Create 3D array for microscopic cross sections
         microxs_arr = np.zeros((len(nuclides), len(mts), 1))
 
-        # Create a material with all nuclides
-        mat_all_nucs = openmc.Material()
-        for nuc in nuclides:
-            if nuc in nuclides_with_data:
-                mat_all_nucs.add_nuclide(nuc, 1.0)
-        mat_all_nucs.set_density("atom/b-cm", 1.0)
-
-        # Create simple model containing the above material
-        surf1 = openmc.Sphere(boundary_type="vacuum")
-        surf1_cell = openmc.Cell(fill=mat_all_nucs, region=-surf1)
-        model = openmc.Model()
-        model.geometry = openmc.Geometry([surf1_cell])
-        model.settings = openmc.Settings(
-            particles=1, batches=1, output={'summary': False})
-
-        with change_directory(tmpdir=True):
-            # Export model within temporary directory
-            model.export_to_model_xml()
-
-            with openmc.lib.run_in_memory(**init_kwargs):
-                # For each nuclide and reaction, compute the flux-averaged
-                # cross section
-                for nuc_index, nuc in enumerate(nuclides):
-                    if nuc not in nuclides_with_data:
-                        continue
-                    lib_nuc = openmc.lib.nuclides[nuc]
-                    for mt_index, mt in enumerate(mts):
-                        xs = lib_nuc.collapse_rate(
-                            mt, temperature, energies, multigroup_flux
-                        )
-                        microxs_arr[nuc_index, mt_index, 0] = xs
+        def compute_microxs():
+            # For each nuclide and reaction, compute the flux-averaged xs
+            for nuc_index, nuc in enumerate(nuclides):
+                if nuc not in nuclides_with_data:
+                    continue
+                lib_nuc = openmc.lib.load_nuclide(nuc)
+                for mt_index, mt in enumerate(mts):
+                    microxs_arr[nuc_index, mt_index, 0] = lib_nuc.collapse_rate(
+                        mt, temperature, energies, multigroup_flux
+                    )
+
+        # Compute microscopic cross sections within a temporary session
+        if not openmc.lib.is_initialized:
+            with openmc.lib.TemporarySession(**init_kwargs):
+                compute_microxs()
+        else:
+            compute_microxs()
 
         return cls(microxs_arr, nuclides, reactions)
 

openmc/lib/core.py

@@ -4,7 +4,9 @@
                     c_uint64, c_size_t)
 import sys
 import os
+from pathlib import Path
 from random import getrandbits
+from tempfile import TemporaryDirectory
 
 import numpy as np
 from numpy.ctypeslib import as_array
@@ -617,6 +619,69 @@ def run_in_memory(**kwargs):
         finalize()
 
 
+class TemporarySession:
+    """Context manager for running via openmc.lib in a temporary directory.
+
+    This class is useful for accessing functionality from openmc.lib without
+    polluting your current working directory with OpenMC files. It is used
+    internally as a persistent session to avoid loading cross sections multiple
+    times.
+
+    Parameters
+    ----------
+    model : openmc.Model, optional
+        OpenMC model to use for the session. If None, a minimal working model is
+        created.
+    **init_kwargs
+        Keyword arguments to pass to :func:`openmc.lib.init`.
+
+    Attributes
+    ----------
+    model : openmc.Model
+        The OpenMC model used for the session.
+
+    """
+    def __init__(self, model=None, **init_kwargs):
+        self.init_kwargs = init_kwargs
+        if model is None:
+            surf = openmc.Sphere(boundary_type="vacuum")
+            cell = openmc.Cell(region=-surf)
+            model = openmc.Model()
+            model.geometry = openmc.Geometry([cell])
+            model.settings = openmc.Settings(
+                particles=1, batches=1, output={'summary': False})
+        self.model = model
+
+    def __enter__(self):
+        """Initialize the OpenMC library in a temporary directory."""
+        # Make sure OpenMC is not already initialized
+        if openmc.lib.is_initialized:
+            raise RuntimeError("openmc.lib is already initialized.")
+
+        # Store original working directory
+        self.orig_dir = Path.cwd()
+
+        # Set up temporary directory
+        self.tmp_dir = TemporaryDirectory()
+        working_dir = Path(self.tmp_dir.name)
+        working_dir.mkdir(parents=True, exist_ok=True)
+        os.chdir(working_dir)
+
+        # Export model and initialize OpenMC
+        self.model.export_to_model_xml()
+        openmc.lib.init(**self.init_kwargs)
+
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Finalize the OpenMC library and clean up temporary directory."""
+        try:
+            openmc.lib.finalize()
+        finally:
+            os.chdir(self.orig_dir)
+            self.tmp_dir.cleanup()
+
+
 class _DLLGlobal:
     """Data descriptor that exposes global variables from libopenmc."""
     def __init__(self, ctype, name):

openmc/lib/nuclide.py

@@ -40,8 +40,14 @@ def load_nuclide(name):
     name : str
         Name of the nuclide, e.g. 'U235'
 
+    Returns
+    -------
+    Nuclide
+        The class:`Nuclide` that was just loaded.
+
     """
     _dll.openmc_load_nuclide(name.encode(), None, 0)
+    return nuclides[name]
 
 
 class Nuclide(_FortranObject):

openmc/mesh.py

@@ -400,31 +400,30 @@ def material_volumes(
         """
         import openmc.lib
 
-        with change_directory(tmpdir=True):
-            # In order to get mesh into model, we temporarily replace the
-            # tallies with a single mesh tally using the current mesh
-            original_tallies = model.tallies
-            new_tally = openmc.Tally()
-            new_tally.filters = [openmc.MeshFilter(self)]
-            new_tally.scores = ['flux']
-            model.tallies = [new_tally]
-
-            # Export model to XML
-            model.export_to_model_xml()
-
-            # Get material volume fractions
-            kwargs.setdefault('output', True)
-            if 'args' in kwargs:
-                kwargs['args'] = ['-c'] + kwargs['args']
-            kwargs.setdefault('args', ['-c'])
-            openmc.lib.init(**kwargs)
+        # In order to get mesh into model, we temporarily replace the
+        # tallies with a single mesh tally using the current mesh
+        original_tallies = model.tallies
+        new_tally = openmc.Tally()
+        new_tally.filters = [openmc.MeshFilter(self)]
+        new_tally.scores = ['flux']
+        model.tallies = [new_tally]
+
+        # Set default arguments
+        kwargs.setdefault('output', True)
+        if 'args' in kwargs:
+            kwargs['args'] = ['-c'] + kwargs['args']
+        kwargs.setdefault('args', ['-c'])
+
+        with openmc.lib.TemporarySession(model, **kwargs):
+            # Get mesh from single tally
             mesh = openmc.lib.tallies[new_tally.id].filters[0].mesh
+
+            # Compute material volumes
             volumes = mesh.material_volumes(
                 n_samples, max_materials, output=kwargs['output'])
-            openmc.lib.finalize()
 
-            # Restore original tallies
-            model.tallies = original_tallies
+        # Restore original tallies
+        model.tallies = original_tallies
 
         return volumes
 

openmc/weight_windows.py

@@ -1041,10 +1041,6 @@ def export_to_hdf5(self, path: PathLike = 'weight_windows.h5', **init_kwargs):
         # Get absolute path before moving to temporary directory
         path = Path(path).resolve()
 
-        with change_directory(tmpdir=True):
-            # Write the model to an XML file
-            model.export_to_model_xml()
-
-            # Load the model with openmc.lib and then export it to an HDF5 file
-            with openmc.lib.run_in_memory(**init_kwargs):
-                openmc.lib.export_weight_windows(path)
+        # Load the model with openmc.lib and then export it to an HDF5 file
+        with openmc.lib.TemporarySession(model, **init_kwargs):
+            openmc.lib.export_weight_windows(path)

src/nuclide.cpp

@@ -1114,7 +1114,7 @@ extern "C" size_t nuclides_size()
 extern "C" int openmc_load_nuclide(const char* name, const double* temps, int n)
 {
   if (data::nuclide_map.find(name) == data::nuclide_map.end() ||
-      data::nuclide_map.at(name) >= data::elements.size()) {
+      data::nuclide_map.at(name) >= data::nuclides.size()) {
     LibraryKey key {Library::Type::neutron, name};
     const auto& it = data::library_map.find(key);
     if (it == data::library_map.end()) {
@@ -1215,7 +1215,6 @@ extern "C" int openmc_nuclide_collapse_rate(int index, int MT,
     *xs = data::nuclides[index]->collapse_rate(
       MT, temperature, {energy, energy + n + 1}, {flux, flux + n});
   } catch (const std::out_of_range& e) {
-    fmt::print("Caught error\n");
     set_errmsg(e.what());
     return OPENMC_E_OUT_OF_BOUNDS;
   }

tests/unit_tests/test_deplete_chain.py

@@ -310,8 +310,7 @@ def test_capture_branch_infer_ground():
     # Create nuclide to be added into the chain
     xe136m = nuclide.Nuclide("Xe136_m1")
 
-    chain.nuclides.append(xe136m)
-    chain.nuclide_dict[xe136m.name] = len(chain.nuclides) - 1
+    chain.add_nuclide(xe136m)
 
     chain.set_branch_ratios(infer_br, "(n,gamma)")
 
@@ -327,8 +326,7 @@ def test_capture_branch_no_rxn():
 
     u5m = nuclide.Nuclide("U235_m1")
 
-    chain.nuclides.append(u5m)
-    chain.nuclide_dict[u5m.name] = len(chain.nuclides) - 1
+    chain.add_nuclide(u5m)
 
     with pytest.raises(AttributeError, match="U234"):
         chain.set_branch_ratios(u4br)