Don't write reaction rates in depletion results by default, remove per-stage data for multistage integrators (#3609)

Author: paulromano

docs/source/io_formats/depletion_results.rst

@@ -4,30 +4,28 @@
 Depletion Results File Format
 =============================
 
-The current version of the depletion results file format is 1.1.
+The current version of the depletion results file format is 1.2.
 
 **/**
 
 :Attributes: - **filetype** (*char[]*) -- String indicating the type of file.
              - **version** (*int[2]*) -- Major and minor version of the
                statepoint file format.
 
-:Datasets: - **eigenvalues** (*double[][][2]*) -- k-eigenvalues at each
-             time/stage. This array has shape (number of timesteps, number of
-             stages, value). The last axis contains the eigenvalue and the
-             associated uncertainty
-           - **number** (*double[][][][]*) -- Total number of atoms. This array
-             has shape (number of timesteps, number of stages, number of
+:Datasets: - **eigenvalues** (*double[][2]*) -- k-eigenvalues at each timestep.
+             This array has shape (number of timesteps, 2). The second axis
+             contains the eigenvalue and its associated uncertainty.
+           - **number** (*double[][][]*) -- Total number of atoms at each
+             timestep. This array has shape (number of timesteps, number of
              materials, number of nuclides).
-           - **reaction rates** (*double[][][][][]*) -- Reaction rates used to
-             build depletion matrices. This array has shape (number of
-             timesteps, number of stages, number of materials, number of
-             nuclides, number of reactions).
+           - **reaction rates** (*double[][][][]*) -- Reaction rates at each
+             timestep. This array has shape (number of timesteps, number of
+             materials, number of nuclides, number of reactions). Only stored if
+             write_rates=True.
            - **time** (*double[][2]*) -- Time in [s] at beginning/end of each
              step.
-           - **source_rate** (*double[][]*) -- Power in [W] or source rate in
-             [neutron/sec]. This array has shape (number of timesteps, number
-             of stages).
+           - **source_rate** (*double[]*) -- Power in [W] or source rate in
+             [neutron/sec] for each timestep.
            - **depletion time** (*double[]*) -- Average process time in [s]
              spent depleting a material across all burnable materials and,
              if applicable, MPI processes.

openmc/deplete/abc.py

@@ -631,17 +631,7 @@ def __init__(
             solver: str = "cram48",
             continue_timesteps: bool = False,
         ):
-        # Check number of stages previously used
-        if operator.prev_res is not None:
-            res = operator.prev_res[-1]
-            if res.data.shape[0] != self._num_stages:
-                raise ValueError(
-                    "{} incompatible with previous restart calculation. "
-                    "Previous scheme used {} intermediate solutions, while "
-                    "this uses {}".format(
-                        self.__class__.__name__, res.data.shape[0],
-                        self._num_stages))
-        elif continue_timesteps:
+        if continue_timesteps and operator.prev_res is None:
             raise ValueError("Continuation run requires passing prev_results.")
         self.operator = operator
         self.chain = operator.chain
@@ -775,12 +765,8 @@ def __call__(
         -------
         proc_time : float
             Time spent in CRAM routines for all materials in [s]
-        n_list : list of list of numpy.ndarray
-            Concentrations at each of the intermediate points with
-            the final concentration as the last element
-        op_results : list of openmc.deplete.OperatorResult
-            Eigenvalue and reaction rates from intermediate transport
-            simulations
+        n_end : list of numpy.ndarray
+            Concentrations at end of timestep
         """
 
     @property
@@ -811,9 +797,9 @@ def _get_bos_data_from_restart(self, source_rate, bos_conc):
         """Get beginning of step concentrations, reaction rates from restart"""
         res = self.operator.prev_res[-1]
         # Depletion methods expect list of arrays
-        bos_conc = list(res.data[0])
-        rates = res.rates[0]
-        k = ufloat(res.k[0, 0], res.k[0, 1])
+        bos_conc = list(res.data)
+        rates = res.rates
+        k = ufloat(res.k[0], res.k[1])
 
         if res.source_rate != 0.0:
             # Scale reaction rates by ratio of source rates
@@ -855,7 +841,8 @@ def integrate(
             self,
             final_step: bool = True,
             output: bool = True,
-            path: PathLike = 'depletion_results.h5'
+            path: PathLike = 'depletion_results.h5',
+            write_rates: bool = False
         ):
         """Perform the entire depletion process across all steps
 
@@ -874,6 +861,11 @@ def integrate(
             Path to file to write. Defaults to 'depletion_results.h5'.
 
             .. versionadded:: 0.15.0
+        write_rates : bool, optional
+            Whether reaction rates should be written to the results file for
+            each step. Defaults to ``False`` to reduce file size.
+
+            .. versionadded:: 0.15.3
         """
         with change_directory(self.operator.output_dir):
             n = self.operator.initial_condition()
@@ -890,18 +882,22 @@ def integrate(
                     n, res = self._get_bos_data_from_restart(source_rate, n)
 
                 # Solve Bateman equations over time interval
-                proc_time, n_list, res_list = self(n, res.rates, dt, source_rate, i)
-
-                # Insert BOS concentration, transport results
-                n_list.insert(0, n)
-                res_list.insert(0, res)
-
-                # Remove actual EOS concentration for next step
-                n = n_list.pop()
-
-                StepResult.save(self.operator, n_list, res_list, [t, t + dt],
-                                source_rate, self._i_res + i, proc_time, path)
+                proc_time, n_end = self(n, res.rates, dt, source_rate, i)
+
+                StepResult.save(
+                    self.operator,
+                    n,
+                    res,
+                    [t, t + dt],
+                    source_rate,
+                    self._i_res + i,
+                    proc_time,
+                    write_rates=write_rates,
+                    path=path
+                )
 
+                # Update for next step
+                n = n_end
                 t += dt
 
             # Final simulation -- in the case that final_step is False, a zero
@@ -910,9 +906,18 @@ def integrate(
             # solve)
             if output and final_step and comm.rank == 0:
                 print(f"[openmc.deplete] t={t} (final operator evaluation)")
-            res_list = [self.operator(n, source_rate if final_step else 0.0)]
-            StepResult.save(self.operator, [n], res_list, [t, t],
-                         source_rate, self._i_res + len(self), proc_time, path)
+            res_final = self.operator(n, source_rate if final_step else 0.0)
+            StepResult.save(
+                self.operator,
+                n,
+                res_final,
+                [t, t],
+                source_rate,
+                self._i_res + len(self),
+                proc_time,
+                write_rates=write_rates,
+                path=path
+            )
             self.operator.write_bos_data(len(self) + self._i_res)
 
         self.operator.finalize()
@@ -1171,10 +1176,40 @@ def _get_bos_data_from_operator(self, step_index, step_power, n_bos):
             self.operator.settings.particles //= self.n_steps
         return inherited
 
+    @abstractmethod
+    def __call__(self, n, rates, dt, source_rate, i):
+        """Perform the integration across one time step
+
+        Parameters
+        ----------
+        n : list of numpy.ndarray
+            List of atom number arrays for each material. Each array has
+            shape ``(n_nucs,)`` where ``n_nucs`` is the number of nuclides
+        rates : openmc.deplete.ReactionRates
+            Reaction rates (from transport operator)
+        dt : float
+            Time step in [s]
+        source_rate : float
+            Power in [W] or source rate in [neutron/sec]
+        i : int
+            Current time step index
+
+        Returns
+        -------
+        proc_time : float
+            Time spent in transport simulation
+        n_end : list of numpy.ndarray
+            Updated atom number densities for each material
+        op_result : OperatorResult
+            Eigenvalue and reaction rates resulting from transport simulation
+
+        """
+
     def integrate(
         self,
         output: bool = True,
-        path: PathLike = "depletion_results.h5"
+        path: PathLike = "depletion_results.h5",
+        write_rates: bool = False
     ):
         """Perform the entire depletion process across all steps
 
@@ -1186,11 +1221,17 @@ def integrate(
             Path to file to write. Defaults to 'depletion_results.h5'.
 
             .. versionadded:: 0.15.0
+        write_rates : bool, optional
+            Whether reaction rates should be written to the results file for
+            each step. Defaults to ``False`` to reduce file size.
+
+            .. versionadded:: 0.15.3
         """
         with change_directory(self.operator.output_dir):
             n = self.operator.initial_condition()
             t, self._i_res = self._get_start_data()
 
+            res_end = None  # Will be set in first iteration
             for i, (dt, p) in enumerate(self):
                 if output:
                     print(f"[openmc.deplete] t={t} s, dt={dt} s, source={p}")
@@ -1200,28 +1241,38 @@ def integrate(
                         n, res = self._get_bos_data_from_operator(i, p, n)
                     else:
                         n, res = self._get_bos_data_from_restart(p, n)
-                else:
-                    # Pull rates, k from previous iteration w/o
-                    # re-running transport
-                    res = res_list[-1]  # defined in previous i iteration
-
-                proc_time, n_list, res_list = self(n, res.rates, dt, p, i)
 
-                # Insert BOS concentration, transport results
-                n_list.insert(0, n)
-                res_list.insert(0, res)
-
-                # Remove actual EOS concentration for next step
-                n = n_list.pop()
-
-                StepResult.save(self.operator, n_list, res_list, [t, t + dt],
-                             p, self._i_res + i, proc_time, path)
+                proc_time, n_end, res_end = self(n, res.rates, dt, p, i)
+
+                StepResult.save(
+                    self.operator,
+                    n,
+                    res,
+                    [t, t + dt],
+                    p,
+                    self._i_res + i,
+                    proc_time,
+                    write_rates=write_rates,
+                    path=path
+                )
 
+                # Update for next step
+                n = n_end
+                res = res_end
                 t += dt
 
             # No final simulation for SIE, use last iteration results
-            StepResult.save(self.operator, [n], [res_list[-1]], [t, t],
-                         p, self._i_res + len(self), proc_time, path)
+            StepResult.save(
+                self.operator,
+                n,
+                res_end,
+                [t, t],
+                p,
+                self._i_res + len(self),
+                proc_time,
+                write_rates=write_rates,
+                path=path
+            )
             self.operator.write_bos_data(self._i_res + len(self))
 
         self.operator.finalize()