Allow already-initialized openmc.lib in TemporarySession (#3505)

openmc/deplete/microxs.py

@@ -352,7 +352,8 @@ class MicroXS:
         # Create 3D array for microscopic cross sections
         microxs_arr = np.zeros((len(nuclides), len(mts), 1))
 
-        def compute_microxs():
+        # Compute microscopic cross sections within a temporary session
+        with openmc.lib.TemporarySession(**init_kwargs):
             # For each nuclide and reaction, compute the flux-averaged xs
             for nuc_index, nuc in enumerate(nuclides):
                 if nuc not in nuclides_with_data:
@@ -363,13 +364,6 @@ class MicroXS:
                         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)
 
     @classmethod

openmc/lib/core.py

@@ -654,9 +654,10 @@ class TemporarySession:
 
     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.")
+        # If already initialized, the context manager is a no-op
+        self.already_initialized = openmc.lib.is_initialized
+        if self.already_initialized:
+            return self
 
         # Store original working directory
         self.orig_dir = Path.cwd()
@@ -675,8 +676,11 @@ class TemporarySession:
 
     def __exit__(self, exc_type, exc_value, traceback):
         """Finalize the OpenMC library and clean up temporary directory."""
+        if self.already_initialized:
+            return
+
         try:
-            openmc.lib.finalize()
+            finalize()
         finally:
             os.chdir(self.orig_dir)
             self.tmp_dir.cleanup()

openmc/model/model.py

@@ -996,16 +996,13 @@ class Model:
         plot_obj.v_res = pixels[1]
         plot_obj.basis = basis
 
-        if self.is_initialized:
-            return openmc.lib.id_map(plot_obj)
-        else:
-            # Silence output by default. Also set arguments to start in volume
-            # calculation mode to avoid loading cross sections
-            init_kwargs.setdefault('output', False)
-            init_kwargs.setdefault('args', ['-c'])
+        # Silence output by default. Also set arguments to start in volume
+        # calculation mode to avoid loading cross sections
+        init_kwargs.setdefault('output', False)
+        init_kwargs.setdefault('args', ['-c'])
 
-            with openmc.lib.TemporarySession(self, **init_kwargs):
-                return openmc.lib.id_map(plot_obj)
+        with openmc.lib.TemporarySession(self, **init_kwargs):
+            return openmc.lib.id_map(plot_obj)
 
     @add_plot_params
     def plot(
@@ -1229,20 +1226,15 @@ class Model:
         """
         import openmc.lib
 
-        if self.is_initialized:
+        # Silence output by default. Also set arguments to start in volume
+        # calculation mode to avoid loading cross sections
+        init_kwargs.setdefault('output', False)
+        init_kwargs.setdefault('args', ['-c'])
+
+        with openmc.lib.TemporarySession(self, **init_kwargs):
             return openmc.lib.sample_external_source(
                 n_samples=n_samples, prn_seed=prn_seed
             )
-        else:
-            # Silence output by default. Also set arguments to start in volume
-            # calculation mode to avoid loading cross sections
-            init_kwargs.setdefault('output', False)
-            init_kwargs.setdefault('args', ['-c'])
-
-            with openmc.lib.TemporarySession(self, **init_kwargs):
-                return openmc.lib.sample_external_source(
-                    n_samples=n_samples, prn_seed=prn_seed
-                )
 
     def apply_tally_results(self, statepoint: PathLike | openmc.StatePoint):
         """Apply results from a statepoint to tally objects on the Model