Fix in ReactionRates class. More documentation updates

openmc/deplete/__init__.py

@@ -16,7 +16,7 @@ except ImportError:
 
 from .nuclide import *
 from .chain import *
-from .openmc_wrapper import *
+from .operator import *
 from .reaction_rates import *
 from .abc import *
 from .results import *

openmc/deplete/abc.py

@@ -66,7 +66,7 @@ class TransportOperator(metaclass=ABCMeta):
 
         Parameters
         ----------
-        vec : list of numpy.array
+        vec : list of numpy.ndarray
             Total atoms to be used in function.
         print_out : bool, optional
             Whether or not to print out time.
@@ -108,7 +108,7 @@ class TransportOperator(metaclass=ABCMeta):
 
         Returns
         -------
-        list of numpy.array
+        list of numpy.ndarray
             Total density for initial conditions.
         """
 

openmc/deplete/atom_number.py

@@ -55,7 +55,7 @@ class AtomNumber(object):
 
         self.n_nuc_burn = n_nuc_burn
 
-        self.number = np.zeros((len(local_mats), self.n_nuc))
+        self.number = np.zeros((len(local_mats), len(nuclides)))
 
     def __getitem__(self, pos):
         """Retrieves total atom number from AtomNumber.
@@ -69,7 +69,7 @@ class AtomNumber(object):
 
         Returns
         -------
-        numpy.array
+        numpy.ndarray
             The value indexed from self.number.
         """
 
@@ -153,7 +153,7 @@ class AtomNumber(object):
             Material index.
         nuc : str, int or slice
             Nuclide index.
-        val : numpy.array
+        val : numpy.ndarray
             Array of densities to set in [atom/cm^3]
 
         """
@@ -190,7 +190,7 @@ class AtomNumber(object):
         ----------
         mat : str, int or slice
             Material index.
-        val : numpy.array
+        val : numpy.ndarray
             The slice to set in [atom]
 
         """

openmc/deplete/operator.py

@@ -11,15 +11,8 @@ import copy
 from collections import OrderedDict
 from itertools import chain
 import os
-import random
-import sys
 import time
-try:
-    import lxml.etree as ET
-    _have_lxml = True
-except ImportError:
-    import xml.etree.ElementTree as ET
-    _have_lxml = False
+import xml.etree.ElementTree as ET
 
 import h5py
 import numpy as np
@@ -34,6 +27,19 @@ from .reaction_rates import ReactionRates
 
 
 def _distribute(items):
+    """Distribute items across MPI communicator
+
+    Parameters
+    ----------
+    items : list
+        List of items of distribute
+
+    Returns
+    -------
+    list
+        Items assigned to process that called
+
+    """
     min_size, extra = divmod(len(items), comm.size)
     j = 0
     for i in range(comm.size):
@@ -114,7 +120,7 @@ class Operator(TransportOperator):
 
         Parameters
         ----------
-        vec : list of numpy.array
+        vec : list of numpy.ndarray
             Total atoms to be used in function.
         power : float
             Power of the reactor in [W]
@@ -241,7 +247,7 @@ class Operator(TransportOperator):
 
         Returns
         -------
-        list of numpy.array
+        list of numpy.ndarray
             Total density for initial conditions.
         """
 

openmc/deplete/reaction_rates.py

@@ -7,14 +7,16 @@ import numpy as np
 
 
 class ReactionRates(np.ndarray):
-    """ReactionRates class.
+    """Reaction rates resulting from a transport operator call
 
-    An ndarray to store reaction rates with string, integer, or slice indexing.
+    This class is a subclass of :class:`numpy.ndarray` with a few custom
+    attributes that make it easy to determine what index corresponds to a given
+    material, nuclide, and reaction rate.
 
     Parameters
     ----------
-    index_mat : OrderedDict of str to int
-        A dictionary mapping material ID as string to index.
+    local_mats : list of str
+        Material IDs
     nuclides : list of str
         Depletable nuclides
     index_rx : OrderedDict of str to int
@@ -36,14 +38,22 @@ class ReactionRates(np.ndarray):
         Number of reactions.
 
     """
-    def __new__(cls, index_mat, nuclides, index_rx):
+
+    # NumPy arrays can be created 1) explicitly 2) using view casting, and 3) by
+    # slicing an existing array. Because of these possibilities, it's necessary
+    # to put initialization logic in __new__ rather than __init__. Additionally,
+    # subclasses need to handle the multiple ways of creating arrays by using
+    # the __array_finalize__ method (discussed here:
+    # https://docs.scipy.org/doc/numpy/user/basics.subclassing.html)
+
+    def __new__(cls, local_mats, nuclides, index_rx):
         # Create appropriately-sized zeroed-out ndarray
-        shape = (len(index_mat), len(nuclides), len(index_rx))
+        shape = (len(local_mats), len(nuclides), len(index_rx))
         obj = super().__new__(cls, shape)
         obj[:] = 0.0
 
         # Add mapping attributes
-        obj.index_mat = index_mat
+        obj.index_mat = {mat: i for i, mat in enumerate(local_mats)}
         obj.index_nuc = {nuc: i for i, nuc in enumerate(nuclides)}
         obj.index_rx = index_rx
 
@@ -56,6 +66,11 @@ class ReactionRates(np.ndarray):
         self.index_nuc = getattr(obj, 'index_nuc', None)
         self.index_rx = getattr(obj, 'index_rx', None)
 
+    # Reaction rates are distributed to other processes via multiprocessing,
+    # which entails pickling the objects. In order to preserve the custom
+    # attributes, we have to modify how the ndarray is pickled as described
+    # here: https://stackoverflow.com/a/26599346/1572453
+
     def __reduce__(self):
         state = super().__reduce__()
         new_state = state[2] + (self.index_mat, self.index_nuc, self.index_rx)
@@ -69,7 +84,7 @@ class ReactionRates(np.ndarray):
 
     @property
     def n_mat(self):
-        """Number of cells."""
+        """Number of materials."""
         return len(self.index_mat)
 
     @property

openmc/deplete/results.py

@@ -42,7 +42,7 @@ class Results(object):
         Number of materials in entire geometry.
     n_stages : int
         Number of stages in simulation.
-    data : numpy.array
+    data : numpy.ndarray
         Atom quantity, stored by stage, mat, then by nuclide.
 
     """

openmc/deplete/utilities.py

@@ -38,6 +38,7 @@ def evaluate_single_nuclide(results, mat, nuc):
 
     return time, concentration
 
+
 def evaluate_reaction_rate(results, mat, nuc, rx):
     """Return reaction rate in a single material/nuclide from a results list.
 

tests/unit_tests/test_deplete_reaction.py

@@ -7,11 +7,11 @@ from openmc.deplete import ReactionRates
 def test_get_set():
     """Tests the get/set methods."""
 
-    mat_to_ind = {"10000" : 0, "10001" : 1}
-    nuc_to_ind = {"U238" : 0, "U235" : 1}
-    react_to_ind = {"fission" : 0, "(n,gamma)" : 1}
+    local_mats = ["10000", "10001"]
+    nuclides = ["U238", "U235"]
+    react_to_ind = {"fission": 0, "(n,gamma)": 1}
 
-    rates = ReactionRates(mat_to_ind, nuc_to_ind, react_to_ind)
+    rates = ReactionRates(local_mats, nuclides, react_to_ind)
     assert rates.shape == (2, 2, 2)
     assert np.all(rates == 0.0)
 
@@ -50,34 +50,14 @@ def test_get_set():
     assert rates.get("10000", "U238", "fission") == 5.0
 
 
-def test_n_mat():
+def test_properties():
     """Test number of materials property."""
-    mat_to_ind = {"10000" : 0, "10001" : 1}
-    nuc_to_ind = {"U238" : 0, "U235" : 1, "Gd157" : 2}
-    react_to_ind = {"fission" : 0, "(n,gamma)" : 1, "(n,2n)" : 2, "(n,3n)" : 3}
+    local_mats = ["10000", "10001"]
+    nuclides = ["U238", "U235", "Gd157"]
+    react_to_ind = {"fission": 0, "(n,gamma)": 1, "(n,2n)": 2, "(n,3n)": 3}
 
-    rates = ReactionRates(mat_to_ind, nuc_to_ind, react_to_ind)
+    rates = ReactionRates(local_mats, nuclides, react_to_ind)
 
     assert rates.n_mat == 2
-
-
-def test_n_nuc():
-    """Test number of nuclides property."""
-    mat_to_ind = {"10000" : 0, "10001" : 1}
-    nuc_to_ind = {"U238" : 0, "U235" : 1, "Gd157" : 2}
-    react_to_ind = {"fission" : 0, "(n,gamma)" : 1, "(n,2n)" : 2, "(n,3n)" : 3}
-
-    rates = ReactionRates(mat_to_ind, nuc_to_ind, react_to_ind)
-
     assert rates.n_nuc == 3
-
-
-def test_n_react():
-    """ Test number of reactions property. """
-    mat_to_ind = {"10000" : 0, "10001" : 1}
-    nuc_to_ind = {"U238" : 0, "U235" : 1, "Gd157" : 2}
-    react_to_ind = {"fission" : 0, "(n,gamma)" : 1, "(n,2n)" : 2, "(n,3n)" : 3}
-
-    rates = ReactionRates(mat_to_ind, nuc_to_ind, react_to_ind)
-
     assert rates.n_react == 4