From e682aa88e4ae5a04779d6388c078e887536e2817 Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Thu, 6 Aug 2020 08:12:36 -0500
Subject: [PATCH] Address @drewejohnson comments on #1628

---
 openmc/deplete/abc.py                       | 15 ++++++++++-----
 tests/unit_tests/test_deplete_activation.py |  2 +-
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/openmc/deplete/abc.py b/openmc/deplete/abc.py
index 7a586a69c9..af32fe3979 100644
--- a/openmc/deplete/abc.py
+++ b/openmc/deplete/abc.py
@@ -633,7 +633,7 @@ class Integrator(ABC):
         Power density of the reactor in [W/gHM]. It is multiplied by
         initial heavy metal inventory to get total power if ``power``
         is not speficied.
-    source_rates : iterable of float
+    source_rates : float or iterable of float, optional
         Source rate in [neutrons/sec] for each interval in :attr:`timesteps`
 
         .. versionadded:: 0.12.1
@@ -934,12 +934,16 @@ class SIIntegrator(Integrator):
         indicates potentially different power levels for each timestep.
         For a 2D problem, the power can be given in [W/cm] as long
         as the "volume" assigned to a depletion material is actually
-        an area in [cm^2]. Either ``power`` or ``power_density`` must be
-        specified.
+        an area in [cm^2]. Either ``power``, ``power_density``, or
+        ``source_rates`` must be specified.
     power_density : float or iterable of float, optional
         Power density of the reactor in [W/gHM]. It is multiplied by
         initial heavy metal inventory to get total power if ``power``
         is not speficied.
+    source_rates : float or iterable of float, optional
+        Source rate in [neutrons/sec] for each interval in :attr:`timesteps`
+
+        .. versionadded:: 0.12.1
     timestep_units : {'s', 'min', 'h', 'd', 'MWd/kg'}
         Units for values specified in the `timesteps` argument. 's' means
         seconds, 'min' means minutes, 'h' means hours, and 'MWd/kg' indicates
@@ -992,11 +996,12 @@ class SIIntegrator(Integrator):
     """
 
     def __init__(self, operator, timesteps, power=None, power_density=None,
-                 timestep_units='s', n_steps=10, solver="cram48"):
+                 source_rates=None, timestep_units='s', n_steps=10,
+                 solver="cram48"):
         check_type("n_steps", n_steps, Integral)
         check_greater_than("n_steps", n_steps, 0)
         super().__init__(
-            operator, timesteps, power, power_density,
+            operator, timesteps, power, power_density, source_rates,
             timestep_units=timestep_units, solver=solver)
         self.n_steps = n_steps
 
diff --git a/tests/unit_tests/test_deplete_activation.py b/tests/unit_tests/test_deplete_activation.py
index d30557f0d4..239ae52502 100644
--- a/tests/unit_tests/test_deplete_activation.py
+++ b/tests/unit_tests/test_deplete_activation.py
@@ -43,7 +43,7 @@ def test_activation(run_in_tmpdir, model):
     sp = model.run()
     with openmc.StatePoint(sp) as sp:
         tally = sp.tallies[1]
-        capture_rate = tally.get_values().ravel()[0]
+        capture_rate = tally.mean.flat[0]
 
     # Create one-nuclide depletion chain
     chain = openmc.deplete.Chain()