Implement support for CellInstanceFilter tally plotting

Author: paulromano

openmc_plotter/docks.py

@@ -406,8 +406,11 @@ def _createFilterTree(self, spatial_filters):
                 continue
 
             def _bin_sort_val(bin):
-                if isinstance(bin, Iterable) and all([isinstance(val, float) for val in bin]):
-                    return np.sum(bin)
+                if isinstance(bin, Iterable):
+                    if all([isinstance(val, float) for val in bin]):
+                        return np.sum(bin)
+                    else:
+                        return tuple(bin)
                 else:
                     return bin
 

openmc_plotter/plotmodel.py

@@ -34,6 +34,7 @@
                     openmc.MaterialFilter,
                     openmc.CellFilter,
                     openmc.DistribcellFilter,
+                    openmc.CellInstanceFilter,
                     openmc.MeshFilter)
 
 _PRODUCTIONS = ('delayed-nu-fission', 'prompt-nu-fission', 'nu-fission',
@@ -330,6 +331,9 @@ def create_tally_image(self, view=None):
 
         units_out = list(units)[0]
 
+        contains_distribcell = tally.contains_filter(openmc.DistribcellFilter)
+        contains_cellinstance = tally.contains_filter(openmc.CellInstanceFilter)
+
         if tally.contains_filter(openmc.MeshFilter):
             if tally_value == 'rel_err':
                 # get both the std. dev. data and mean data
@@ -360,28 +364,26 @@ def create_tally_image(self, view=None):
                                                       nuclides,
                                                       view)
                 return image + (units_out,)
-        elif tally.contains_filter(openmc.DistribcellFilter):
+        elif contains_distribcell or contains_cellinstance:
+            # Select appropriate function
+            if contains_distribcell:
+                create_image = self._create_distribcell_image
+            else:
+                create_image = self._create_cellinstance_image
+
             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)
+                mean_data = create_image(tally, 'mean', scores, nuclides)
+                std_dev_data = create_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)
+                image = create_image(tally, tally_value, scores, nuclides)
                 return image + (units_out,)
         else:
             # same as above, get the std. dev. data
@@ -522,6 +524,26 @@ def _create_distribcell_image(self, tally, tally_value, scores, nuclides):
 
         return image_data, None, data_min, data_max
 
+    def _create_cellinstance_image(self, tally, tally_value, scores, nuclides):
+        cellinst_filter = tally.find_filter(openmc.CellInstanceFilter)
+
+        data = tally.get_values(scores=scores, nuclides=nuclides, value=tally_value)
+        data = data.flatten()
+
+        # create a mask for ids that match the cell
+        image_data = np.full_like(self.ids, np.nan, dtype=float)
+
+        for v, (cell_id, instance) in zip(data, cellinst_filter.bins):
+            cell_id_mask = self.cell_ids == cell_id
+            instance_mask = self.instances == instance
+            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: