only set data of container if writing to frame and necessary

icecube · Sep 6, 2020 · 8a8af11 · 8a8af11
1 parent 18a32eb
commit 8a8af11
Showing 1 changed file with 11 additions and 4 deletions.
diff --git a/ic3_data/data.py b/ic3_data/data.py
@@ -57,6 +57,11 @@ def Configure(self):
         self._dom_exclusions = self._container.config['dom_exclusions']
         self._partial_exclusion = self._container.config['partial_exclusion']
 
+        if self._output_key is None:
+            self._write_to_frame = False
+        else:
+            self._write_to_frame = True
+
         # initalize data fields of data container
         self._container.initialize()
         self._batch_index = 0
@@ -243,7 +248,7 @@ def Physics(self, frame):
         self._container.runtime_batch[self._batch_index] = elapsed_time
 
         # Write data to frame
-        if self._output_key is not None:
+        if self._write_to_frame:
             frame[self._output_key + '_bin_exclusions'] = \
                 self._container.bin_exclusions
             frame[self._output_key + '_bin_indices'] = \
@@ -318,8 +323,9 @@ def _fill_dnn_container(self, om_key, bin_values_list, bin_indices_list,
 
         # Add bin values and indices if non-empty
         if bin_values_list:
-            self._container.bin_values[om_key] = bin_values_list
-            self._container.bin_indices[om_key] = bin_indices_list
+            if self._write_to_frame:
+                self._container.bin_values[om_key] = bin_values_list
+                self._container.bin_indices[om_key] = bin_indices_list
 
             for value, index in zip(bin_values_list, bin_indices_list):
                 if self._is_str_dom_format:
@@ -343,7 +349,8 @@ def _fill_dnn_container(self, om_key, bin_values_list, bin_indices_list,
 
         # Add bin exclusions if non-empty
         if bin_exclusions_list:
-            self._container.bin_exclusions[om_key] = bin_exclusions_list
+            if self._write_to_frame:
+                self._container.bin_exclusions[om_key] = bin_exclusions_list
             for index in bin_exclusions_list:
                 if self._is_str_dom_format:
                     if index == -1: