Repository: openmc-dev/openmc

Issue:
Use of depletion operator as context manager
@yardasol and I were discussing the `TransportOperator` class and how it is used as a context manager by the `Integrator` class:
https://github.com/openmc-dev/openmc/blob/b72c4593a328d9a16f6d3ab1285854527d4c1045/openmc/deplete/abc.py#L778

We both agree that the use of the context manager is not very clear. In fact, most of what the `__enter__` and `__exit__` method are doing is just moving to do a different directory, which doesn't have much to do with the operator at all. If we want to use a context manager, it should just handle moving between directories, and then the call to get the initial condition should be explicit in `Integrator.integrate()`, something like:
```Python
with change_directory(self.operator.output_dir):
    vec = self.operator.initial_condition()
    ...
```

One other point we talked about: the `AtomNumber` class uses the attribute name `n_nuc_burn` which is set to the number of nuclides in the depletion chain, but elsewhere we use "burnable nucs" to mean the intersection of the nuclides in the chain and those that actually have cross section data. To prevent confusion, it would be better to call that variable `n_nuc_chain`.


Discussion:

Resolving patch:
diff --git a/openmc/deplete/abc.py b/openmc/deplete/abc.py
--- a/openmc/deplete/abc.py
+++ b/openmc/deplete/abc.py
@@ -10,6 +10,7 @@
 from copy import deepcopy
 from inspect import signature
 from numbers import Real, Integral
+from contextlib import contextmanager
 import os
 from pathlib import Path
 import sys
@@ -56,6 +57,24 @@
     # Can't set __doc__ on properties on Python 3.4
     pass
 
+@contextmanager
+def change_directory(output_dir):
+    """
+    Helper function for managing the current directory.
+
+    Parameters
+    ----------
+    output_dir : pathlib.Path
+        Directory to switch to.
+    """
+    orig_dir  = os.getcwd()
+    try:
+        output_dir.mkdir(exist_ok=True)
+        os.chdir(output_dir)
+        yield
+    finally:
+        os.chdir(orig_dir)
+
 
 class TransportOperator(ABC):
     """Abstract class defining a transport operator
@@ -87,9 +106,14 @@ class TransportOperator(ABC):
         Initial atom density [atoms/cm^3] to add for nuclides that are zero
         in initial condition to ensure they exist in the decay chain.
         Only done for nuclides with reaction rates.
+    output_dir : pathlib.Path
+        Path to output directory to save results.
     prev_res : Results or None
         Results from a previous depletion calculation. ``None`` if no
         results are to be used.
+    chain : openmc.deplete.Chain
+        The depletion chain information necessary to form matrices and tallies.
+
     """
     def __init__(self, chain_file, fission_q=None, dilute_initial=1.0e3,
                  prev_results=None):
@@ -133,18 +157,6 @@ def __call__(self, vec, source_rate):
 
         """
 
-    def __enter__(self):
-        # Save current directory and move to specific output directory
-        self._orig_dir = os.getcwd()
-        self.output_dir.mkdir(exist_ok=True)
-        os.chdir(self.output_dir)
-
-        return self.initial_condition()
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        self.finalize()
-        os.chdir(self._orig_dir)
-
     @property
     def output_dir(self):
         return self._output_dir
@@ -775,7 +787,8 @@ def integrate(self, final_step=True, output=True):
 
             .. versionadded:: 0.13.1
         """
-        with self.operator as conc:
+        with change_directory(self.operator.output_dir):
+            conc = self.operator.initial_condition()
             t, self._i_res = self._get_start_data()
 
             for i, (dt, source_rate) in enumerate(self):
@@ -814,6 +827,8 @@ def integrate(self, final_step=True, output=True):
                          source_rate, self._i_res + len(self), proc_time)
             self.operator.write_bos_data(len(self) + self._i_res)
 
+        self.operator.finalize()
+
 
 @add_params
 class SIIntegrator(Integrator):
@@ -931,7 +946,8 @@ def integrate(self, output=True):
 
             .. versionadded:: 0.13.1
         """
-        with self.operator as conc:
+        with change_directory(self.operator.output_dir):
+            conc = self.operator.initial_condition()
             t, self._i_res = self._get_start_data()
 
             for i, (dt, p) in enumerate(self):
@@ -967,6 +983,8 @@ def integrate(self, output=True):
                          p, self._i_res + len(self), proc_time)
             self.operator.write_bos_data(self._i_res + len(self))
 
+        self.operator.finalize()
+
 
 class DepSystemSolver(ABC):
     r"""Abstract class for solving depletion equations


Test patch: