Merge pull request #528 from wbinventor/sparse-tallies

Sparse Tallies

openmc/mgxs/library.py

@@ -73,6 +73,9 @@ class Library(object):
     name : str, optional
         Name of the multi-group cross section library. Used as a label to
         identify tallies in OpenMC 'tallies.xml' file.
+    sparse : bool
+        Whether or not the Library's tallies use SciPy's LIL sparse matrix
+        format for compressed data storage
 
     """
 
@@ -92,6 +95,7 @@ def __init__(self, openmc_geometry, by_nuclide=False,
         self._all_mgxs = OrderedDict()
         self._sp_filename = None
         self._keff = None
+        self._sparse = False
 
         self.name = name
         self.openmc_geometry = openmc_geometry
@@ -119,6 +123,7 @@ def __deepcopy__(self, memo):
             clone._all_mgxs = self.all_mgxs
             clone._sp_filename = self._sp_filename
             clone._keff = self._keff
+            clone._sparse = self.sparse
 
             clone._all_mgxs = OrderedDict()
             for domain in self.domains:
@@ -208,6 +213,10 @@ def sp_filename(self):
     def keff(self):
         return self._keff
 
+    @property
+    def sparse(self):
+        return self._sparse
+
     @openmc_geometry.setter
     def openmc_geometry(self, openmc_geometry):
         cv.check_type('openmc_geometry', openmc_geometry, openmc.Geometry)
@@ -286,6 +295,28 @@ def tally_trigger(self, tally_trigger):
         cv.check_type('tally trigger', tally_trigger, openmc.Trigger)
         self._tally_trigger = tally_trigger
 
+    @sparse.setter
+    def sparse(self, sparse):
+        """Convert tally data from NumPy arrays to SciPy list of lists (LIL)
+        sparse matrices, and vice versa.
+
+        This property may be used to reduce the amount of data in memory during
+        tally data processing. The tally data will be stored as SciPy LIL
+        matrices internally within the Tally object. All tally data access
+        properties and methods will return data as a dense NumPy array.
+
+        """
+
+        cv.check_type('sparse', sparse, bool)
+
+        # Sparsify or densify each MGXS in the Library
+        for domain in self.domains:
+            for mgxs_type in self.mgxs_types:
+                mgxs = self.get_mgxs(domain, mgxs_type)
+                mgxs.sparse = self.sparse
+
+        self._sparse = sparse
+
     def build_library(self):
         """Initialize MGXS objects in each domain and for each reaction type
         in the library.
@@ -381,6 +412,7 @@ def load_from_statepoint(self, statepoint):
             for mgxs_type in self.mgxs_types:
                 mgxs = self.get_mgxs(domain, mgxs_type)
                 mgxs.load_from_statepoint(statepoint)
+                mgxs.sparse = self.sparse
 
     def get_mgxs(self, domain, mgxs_type):
         """Return the MGXS object for some domain and reaction rate type.

openmc/mgxs/mgxs.py

@@ -103,6 +103,9 @@ class MGXS(object):
     nuclides : list of str or 'sum'
         A list of nuclide string names (e.g., 'U-238', 'O-16') when by_nuclide
         is True and 'sum' when by_nuclide is False.
+    sparse : bool
+        Whether or not the MGXS' tallies use SciPy's LIL sparse matrix format
+        for compressed data storage
 
     """
 
@@ -121,6 +124,7 @@ def __init__(self, domain=None, domain_type=None,
         self._tally_trigger = None
         self._tallies = None
         self._xs_tally = None
+        self._sparse = False
 
         self.name = name
         self.by_nuclide = by_nuclide
@@ -146,6 +150,7 @@ def __deepcopy__(self, memo):
             clone._energy_groups = copy.deepcopy(self.energy_groups, memo)
             clone._tally_trigger = copy.deepcopy(self.tally_trigger, memo)
             clone._xs_tally = copy.deepcopy(self.xs_tally, memo)
+            clone._sparse = self.sparse
 
             clone._tallies = OrderedDict()
             for tally_type, tally in self.tallies.items():
@@ -199,6 +204,10 @@ def tallies(self):
     def xs_tally(self):
         return self._xs_tally
 
+    @property
+    def sparse(self):
+        return self._sparse
+
     @property
     def num_subdomains(self):
         tally = list(self.tallies.values())[0]
@@ -249,6 +258,29 @@ def tally_trigger(self, tally_trigger):
         cv.check_type('tally trigger', tally_trigger, openmc.Trigger)
         self._tally_trigger = tally_trigger
 
+    @sparse.setter
+    def sparse(self, sparse):
+        """Convert tally data from NumPy arrays to SciPy list of lists (LIL)
+        sparse matrices, and vice versa.
+
+        This property may be used to reduce the amount of data in memory during
+        tally data processing. The tally data will be stored as SciPy LIL
+        matrices internally within the Tally object. All tally data access
+        properties and methods will return data as a dense NumPy array.
+
+        """
+
+        cv.check_type('sparse', sparse, bool)
+
+        # Sparsify or densify the derived MGXS tally and its base tallies
+        if self.xs_tally:
+            self.xs_tally.sparse = sparse
+
+        for tally_name in self.tallies:
+                self.tallies[tally_name].sparse = sparse
+
+        self._sparse = sparse
+
     @staticmethod
     def get_mgxs(mgxs_type, domain=None, domain_type=None,
                  energy_groups=None, by_nuclide=False, name=''):
@@ -503,6 +535,7 @@ def compute_xs(self):
         # Remove NaNs which may have resulted from divide-by-zero operations
         self._xs_tally._mean = np.nan_to_num(self.xs_tally.mean)
         self._xs_tally._std_dev = np.nan_to_num(self.xs_tally.std_dev)
+        self.xs_tally.sparse = self.sparse
 
     def load_from_statepoint(self, statepoint):
         """Extracts tallies in an OpenMC StatePoint with the data needed to
@@ -565,6 +598,7 @@ def load_from_statepoint(self, statepoint):
                                             estimator=tally.estimator)
             sp_tally = sp_tally.get_slice(tally.scores, filters,
                                           filter_bins, tally.nuclides)
+            sp_tally.sparse = self.sparse
             self.tallies[tally_type] = sp_tally
 
         # Compute the cross section from the tallies
@@ -716,6 +750,7 @@ def get_condensed_xs(self, coarse_groups):
 
         # Clone this MGXS to initialize the condensed version
         condensed_xs = copy.deepcopy(self)
+        condensed_xs.sparse = False
         condensed_xs.energy_groups = coarse_groups
 
         # Build energy indices to sum across
@@ -754,17 +789,16 @@ def get_condensed_xs(self, coarse_groups):
                     std_dev = np.sqrt(std_dev)
 
             # Reshape condensed data arrays with one dimension for all filters
-            new_shape = \
-                (tally.num_filter_bins, tally.num_nuclides, tally.num_scores,)
-            mean = np.reshape(mean, new_shape)
-            std_dev = np.reshape(std_dev, new_shape)
+            mean = np.reshape(mean, tally.shape)
+            std_dev = np.reshape(std_dev, tally.shape)
 
             # Override tally's data with the new condensed data
             tally._mean = mean
             tally._std_dev = std_dev
 
         # Compute the energy condensed multi-group cross section
         condensed_xs.compute_xs()
+        condensed_xs.sparse = self.sparse
         return condensed_xs
 
     def get_subdomain_avg_xs(self, subdomains='all'):
@@ -807,8 +841,9 @@ def get_subdomain_avg_xs(self, subdomains='all'):
 
         # Clone this MGXS to initialize the subdomain-averaged version
         avg_xs = copy.deepcopy(self)
+        avg_xs.sparse = False
 
-            # If domain is distribcell, make the new domain 'cell'
+        # If domain is distribcell, make the new domain 'cell'
         if self.domain_type == 'distribcell':
             avg_xs.domain_type = 'cell'
 
@@ -836,18 +871,16 @@ def get_subdomain_avg_xs(self, subdomains='all'):
                 domain_filter.num_bins = 1
 
             # Reshape averaged data arrays with one dimension for all filters
-            new_shape = \
-                (tally.num_filter_bins, tally.num_nuclides, tally.num_scores,)
-            mean = np.reshape(mean, new_shape)
-            std_dev = np.reshape(std_dev, new_shape)
+            mean = np.reshape(mean, tally.shape)
+            std_dev = np.reshape(std_dev, tally.shape)
 
             # Override tally's data with the new condensed data
             tally._mean = mean
             tally._std_dev = std_dev
 
         # Compute the subdomain-averaged multi-group cross section
         avg_xs.compute_xs()
-
+        avg_xs.sparse = self.sparse
         return avg_xs
 
     def print_xs(self, subdomains='all', nuclides='all', xs_type='macro'):

openmc/opencg_compatible.py

@@ -901,7 +901,9 @@ def get_opencg_lattice(openmc_lattice):
 
     # Convert 2D universes array to 3D for OpenCG
     if len(universes.shape) == 2:
-        universes.shape = (1,) + universes.shape
+        new_universes = universes.copy()
+        new_universes.shape = (1,) + universes.shape
+        universes = new_universes
 
     # Initialize an empty array for the OpenCG nested Universes in this Lattice
     universe_array = np.ndarray(tuple(np.array(dimension)[::-1]),

openmc/statepoint.py

@@ -75,6 +75,9 @@ class StatePoint(object):
         energy of the source site.
     source_present : bool
         Indicate whether source sites are present
+    sparse : bool
+        Whether or not the tallies uses SciPy's LIL sparse matrix format for
+        compressed data storage
     tallies : dict
         Dictionary whose keys are tally IDs and whose values are Tally objects
     tallies_present : bool
@@ -110,6 +113,7 @@ def __init__(self, filename):
         self._tallies_read = False
         self._summary = False
         self._global_tallies = None
+        self._sparse = False
 
     def close(self):
         self._f.close()
@@ -318,6 +322,10 @@ def source(self):
     def source_present(self):
         return self._f['source_present'].value > 0
 
+    @property
+    def sparse(self):
+        return self._sparse
+
     @property
     def tallies(self):
         if not self._tallies_read:
@@ -340,7 +348,8 @@ def tallies(self):
             for tally_key in tally_keys:
 
                 # Read the Tally size specifications
-                n_realizations = self._f['{0}{1}/n_realizations'.format(base, tally_key)].value
+                n_realizations = \
+                    self._f['{0}{1}/n_realizations'.format(base, tally_key)].value
 
                 # Create Tally object and assign basic properties
                 tally = openmc.Tally(tally_id=tally_key)
@@ -350,15 +359,17 @@ def tallies(self):
                 tally.num_realizations = n_realizations
 
                 # Read the number of Filters
-                n_filters = self._f['{0}{1}/n_filters'.format(base, tally_key)].value
+                n_filters = \
+                    self._f['{0}{1}/n_filters'.format(base, tally_key)].value
 
                 subbase = '{0}{1}/filter '.format(base, tally_key)
 
                 # Initialize all Filters
                 for j in range(1, n_filters+1):
 
                     # Read the Filter type
-                    filter_type = self._f['{0}{1}/type'.format(subbase, j)].value.decode()
+                    filter_type = \
+                        self._f['{0}{1}/type'.format(subbase, j)].value.decode()
 
                     n_bins = self._f['{0}{1}/n_bins'.format(subbase, j)].value
 
@@ -380,7 +391,8 @@ def tallies(self):
                     tally.add_filter(filter)
 
                 # Read Nuclide bins
-                nuclide_names = self._f['{0}{1}/nuclides'.format(base, tally_key)].value
+                nuclide_names = \
+                    self._f['{0}{1}/nuclides'.format(base, tally_key)].value
 
                 # Add all Nuclides to the Tally
                 for name in nuclide_names:
@@ -415,6 +427,7 @@ def tallies(self):
                     tally.add_score(score)
 
                 # Add Tally to the global dictionary of all Tallies
+                tally.sparse = self.sparse
                 self._tallies[tally_key] = tally
 
             self._tallies_read = True
@@ -439,6 +452,26 @@ def summary(self):
     def with_summary(self):
         return False if self.summary is None else True
 
+    @sparse.setter
+    def sparse(self, sparse):
+        """Convert tally data from NumPy arrays to SciPy list of lists (LIL)
+        sparse matrices, and vice versa.
+
+        This property may be used to reduce the amount of data in memory during
+        tally data processing. The tally data will be stored as SciPy LIL
+        matrices internally within each Tally object. All tally data access
+        properties and methods will return data as a dense NumPy array.
+
+        """
+
+        cv.check_type('sparse', sparse, bool)
+        self._sparse = sparse
+
+        # Update tally sparsities
+        if self._tallies_read:
+            for tally_id in self.tallies:
+                self.tallies[tally_id].sparse = self.sparse
+
     def get_tally(self, scores=[], filters=[], nuclides=[],
                   name=None, id=None, estimator=None):
         """Finds and returns a Tally object with certain properties.