Merge pull request #87 from pshriwise/distribcell

Adding support for viewing distribcell tallies

Author: paulromano

openmc_plotter/plotmodel.py

@@ -33,6 +33,7 @@
 _SPATIAL_FILTERS = (openmc.UniverseFilter,
                     openmc.MaterialFilter,
                     openmc.CellFilter,
+                    openmc.DistribcellFilter,
                     openmc.MeshFilter)
 
 _PRODUCTIONS = ('delayed-nu-fission', 'prompt-nu-fission', 'nu-fission',
@@ -280,6 +281,19 @@ def storeCurrent(self):
         self.previousViews.append(copy.deepcopy(self.currentView))
 
     def create_tally_image(self, view=None):
+        """
+        Parameters
+        ----------
+        view :
+            View used to set bounds of the tally data
+
+        Returns
+        -------
+        tuple
+            image data (numpy.ndarray), data extents (optional),
+            data_min_value (float), data_max_value (float),
+            data label (str)
+        """
         if view is None:
             view = self.currentView
 
@@ -346,20 +360,41 @@ def create_tally_image(self, view=None):
                                                       nuclides,
                                                       view)
                 return image + (units_out,)
+        elif tally.contains_filter(openmc.DistribcellFilter):
+            if tally_value == 'rel_err':
+                mean_data = self._create_distribcell_image(tally,
+                                                           'mean',
+                                                           scores,
+                                                           nuclides)
+                std_dev_data = self._create_distribcell_image(tally,
+                                                              'std_dev',
+                                                              scores,
+                                                              nuclides)
+                image_data = 100 * np.divide(std_dev_data[0],
+                                             mean_data[0],
+                                             out=np.zeros_like(mean_data[0]),
+                                             where=mean_data != 0)
+                data_min = np.min(image_data)
+                data_max = np.max(image_data)
+                return image_data, None, data_min, data_max, '% error'
+            else:
+                image = self._create_distribcell_image(tally,
+                                                       tally_value,
+                                                       scores,
+                                                       nuclides)
+                return image + (units_out,)
         else:
             # same as above, get the std. dev. data
             # and mean date to produce the relative error data
             if tally_value == 'rel_err':
                 mean_data = self._create_tally_domain_image(tally,
                                                             'mean',
                                                             scores,
-                                                            nuclides,
-                                                            view)
+                                                            nuclides)
                 std_dev_data = self._create_tally_domain_image(tally,
                                                            'std_dev',
                                                            scores,
-                                                           nuclides,
-                                                           view)
+                                                           nuclides)
                 image_data = 100 * np.divide(std_dev_data[0],
                                              mean_data[0],
                                              out=np.zeros_like(mean_data[0]),
@@ -377,8 +412,7 @@ def create_tally_image(self, view=None):
                 image = self._create_tally_domain_image(tally,
                                                         tally_value,
                                                         scores,
-                                                        nuclides,
-                                                        view)
+                                                        nuclides)
                 return image + (units_out,)
 
     def _create_tally_domain_image(self, tally, tally_value, scores, nuclides, view=None):
@@ -467,10 +501,31 @@ def _do_op(array, tally_value, ax=0):
 
         return image_data, None, data_min, data_max
 
+    def _create_distribcell_image(self, tally, tally_value, scores, nuclides):
+        dfilter = tally.find_filter(openmc.DistribcellFilter)
+
+        data = tally.get_values(scores=scores, nuclides=nuclides, value=tally_value)
+        data = data.flatten()
+
+        cell_id = dfilter.bins[0]
+        # create a mask for ids that match the cell
+        image_data = np.full_like(self.ids, np.nan, dtype=float)
+
+        cell_id_mask = self.cell_ids == cell_id
+        for i, v in enumerate(data):
+            instance_mask = self.instances == i
+            image_data[cell_id_mask & instance_mask] = v
+
+        data_min = np.min(data)
+        data_max = np.max(data)
+        image_data = np.ma.masked_where(image_data < 0.0, image_data)
+
+        return image_data, None, data_min, data_max
+
     def _create_tally_mesh_image(self, tally, tally_value, scores, nuclides, view=None):
         # some variables used throughout
         if view is None:
-            cv = self.currentView
+            view = self.currentView
 
         sp = self.statepoint
         mesh_filter = tally.find_filter(openmc.MeshFilter)