add option to disable addition of min distance value: save disk space

ic3_data/data_formats_detector.py

@@ -188,10 +188,12 @@ def mc_tree_input_data(frame, pulses, config, dom_exclusions,
     distance_bins = [0., 30., 60., 150.]  # meter
     distance_cutoff = 500.  # meter
     energy_cutoff = 1  # GeV
+    add_distance = False
 
     angle_bins_rad = sorted(np.deg2rad(angle_bins))
     distance_bins = sorted(distance_bins)
 
     data_dict = get_mc_tree_input_data_dict(
-        frame, angle_bins_rad, distance_bins, distance_cutoff, energy_cutoff)
+        frame, angle_bins_rad, distance_bins, distance_cutoff, energy_cutoff,
+        add_distance)
     return global_time_offset, data_dict

ic3_data/utils/mc_tree_input.py

@@ -6,7 +6,7 @@
 
 
 def get_mc_tree_input_data_dict(frame, angle_bins, distance_bins,
-                                distance_cutoff, energy_cutoff):
+                                distance_cutoff, energy_cutoff, add_distance):
 
     allowed_types = [dataclasses.I3Particle.NuclInt,
                      dataclasses.I3Particle.PairProd,
@@ -19,11 +19,12 @@ def get_mc_tree_input_data_dict(frame, angle_bins, distance_bins,
 
     num_dist_bins = len(distance_bins) - 1
     num_angle_bins = len(angle_bins) - 1
-    num_bins = 2 + num_dist_bins * num_angle_bins
+    num_bins = 1 + add_distance + num_dist_bins * num_angle_bins
 
     # create empty data array for DOMs
     dom_data = np.zeros(shape=(86, 60, num_bins))
-    dom_data[..., 0] = float('inf')
+    if add_distance:
+        dom_data[..., 0] = float('inf')
 
     # walk through energy losses and calculate data
     for loss in frame['I3MCTree']:
@@ -63,15 +64,17 @@ def get_mc_tree_input_data_dict(frame, angle_bins, distance_bins,
 
                 if not out_of_bounds:
                     # calculate index
-                    index = 1 + index_dist * num_angle_bins + index_angle
-                    assert index > 0 and index < num_bins - 1
+                    index = add_distance + \
+                        index_dist * num_angle_bins + index_angle
+                    assert index >= add_distance and index < num_bins - 1
 
                     # accumulate energy
                     dom_data[om_key.string - 1, om_key.om - 1, index] += \
                         loss.energy
 
                 # check if distance is closest so far
-                if distance < dom_data[om_key.string - 1, om_key.om - 1, 0]:
+                if (add_distance and distance <
+                        dom_data[om_key.string - 1, om_key.om - 1, 0]):
                     dom_data[om_key.string - 1, om_key.om - 1, 0] = distance
 
                 if distance < min_distance:
@@ -103,6 +106,7 @@ def get_mc_tree_input_data_dict(frame, angle_bins, distance_bins,
                     bin_values_compressed.append(v)
                     bin_indices_compressed.append(i)
 
-            data_dict[om_key] = (bin_values_compressed, bin_indices_compressed,
-                                 bin_values, bin_indices_list)
+            if len(bin_values_compressed) > 0:
+                data_dict[om_key] = (bin_values_compressed,
+                                     bin_indices_compressed, bin_exclusions)
     return data_dict

ic3_data_ext/ext_boost.cpp

@@ -267,15 +267,17 @@ inline boost::python::dict get_mc_tree_input_data_dict(
                                   boost::python::object& angle_bins_obj,
                                   boost::python::object& distance_bins_obj,
                                   const double distance_cutoff,
-                                  const double energy_cutoff
+                                  const double energy_cutoff,
+                                  const bool add_distance
                                 ) {
 
     // extract objects
     I3Frame& frame = boost::python::extract<I3Frame&>(frame_obj);
 
     const unsigned int num_dist_bins =  len(distance_bins_obj) - 1;
     const unsigned int num_angle_bins =  len(angle_bins_obj) - 1;
-    const unsigned int num_bins = 2 + num_dist_bins * num_angle_bins;
+    const unsigned int num_bins =
+        1 + add_distance + num_dist_bins * num_angle_bins;
 
     std::vector<double> angle_bins;
     std::vector<double> distance_bins;
@@ -295,14 +297,17 @@ inline boost::python::dict get_mc_tree_input_data_dict(
     const I3MCTree& mctree = frame.Get<I3MCTree>("I3MCTree");
 
     // create data array for DOMs [shape: (86, 60, num_bins)]
+    double init_dist_value = 0.;
+    if (add_distance){
+        init_dist_value = std::numeric_limits<double>::max();
+    }
     std::vector<double>  dom_data[86][60];
     for(size_t string = 1; string <= 86; string++){
         for(size_t dom = 1; dom <= 60; dom++){
             dom_data[string-1][dom-1] = std::vector<double>();
 
             // intialize values
-            dom_data[string-1][dom-1].push_back(
-                                    std::numeric_limits<double>::max());
+            dom_data[string-1][dom-1].push_back(init_dist_value);
             for(size_t i=1; i < num_bins; i++){
                 dom_data[string-1][dom-1].push_back(0);
             }
@@ -366,15 +371,15 @@ inline boost::python::dict get_mc_tree_input_data_dict(
 
                 if (!out_of_bounds){
                     // calculate index
-                    unsigned int index =
-                        1 + index_dist * num_angle_bins + index_angle;
+                    unsigned int index = add_distance +
+                        index_dist * num_angle_bins + index_angle;
 
                     // accumulate energy
                     dom_data[string - 1][dom - 1][index] += loss.GetEnergy();
                 }
 
                 // check if distance is closest so far
-                if (distance < dom_data[string - 1][dom - 1][0]){
+                if (add_distance && distance < dom_data[string - 1][dom - 1][0]){
                     dom_data[string - 1][dom - 1][0] = distance;
                 }