Address some of Chris' comments and lint fixes

XENON1T · Dec 17, 2017 · d61afe0 · d61afe0
1 parent e240fe7
commit d61afe0
Showing 1 changed file with 14 additions and 13 deletions.
diff --git a/hax/treemakers/posrec.py b/hax/treemakers/posrec.py
@@ -8,7 +8,7 @@
 from pax import exceptions
 from pax.plugins.interaction_processing.S1AreaFractionTopProbability import binom_test as binom_test_pax
 from keras.models import model_from_json
-import tensorflow as tf
+
 
 class PositionReconstruction(TreeMaker):
     """Stores position-reconstruction-related variables.
@@ -27,7 +27,7 @@ class PositionReconstruction(TreeMaker):
 
        - x_observed_nn_tf: TensorFlow NN reconstructed x position
        - y_observed_nn_tf: TensorFlow NN reconstructed y position
-       
+
        - s1_area_upper_injection_fraction: s1 area fraction near Rn220 injection points (near PMT 131)
        - s1_area_lower_injection_fraction: s1 area fraction near Rn220 injection points (near PMT 243)
     """
@@ -40,13 +40,13 @@ class PositionReconstruction(TreeMaker):
     def __init__(self):
 
         hax.minitrees.TreeMaker.__init__(self)
-        
+
         # We need to pull some stuff from the pax config
         self.pax_config = load_configuration("XENON1T")
-        self.tpc_channels = list(range(0, 247 + 1))
+        self.tpc_channels = self.pax_config['DEFAULT']['channels_in_detector']['tpc']
+
         self.confused_s1_channels = []
-        self.statistic = self.pax_config['BuildInteractions.BasicInteractionProperties'] \
-                         ['s1_pattern_statistic']
+        self.statistic = self.pax_config['BuildInteractions.BasicInteractionProperties']['s1_pattern_statistic']
         qes = np.array(self.pax_config['DEFAULT']['quantum_efficiencies'])
 
         self.pattern_fitter = PatternFitter(
@@ -78,8 +78,8 @@ def __init__(self):
                               130, 134, 135, 139, 148, 150, 152, 162, 178, 183,
                               190, 198, 206, 213, 214, 234, 239, 244, 27, 73,
                               91, 137, 167, 203]
-        self.ntop_pmts = 127  # How to get this automatically?
-                
+        self.ntop_pmts = len(self.pax_config['DEFAULT']['channels_top'])
+
     def get_data(self, dataset, event_list=None):
 
         data, _ = hax.minitrees.load_single_dataset(dataset, ['Corrections',
@@ -131,14 +131,14 @@ def extract_data(self, event):
         predicted_xy_tensorflow = self.nn_tensorflow.predict(s2apc_clean_norm)
         event_data['x_observed_nn_tf'] = predicted_xy_tensorflow[0, 0] / 10.
         event_data['y_observed_nn_tf'] = predicted_xy_tensorflow[0, 1] / 10.
-        
+
         # s1 area fraction near injection points for Rn220 source
         area_upper_injection = s1.area_per_channel[131] + s1.area_per_channel[138] + s1.area_per_channel[146] + s1.area_per_channel[147]
-        area_lower_injection = s1.area_per_channel[236] + s1.area_per_channel[237] + s1.area_per_channel[243] 
-        
+        area_lower_injection = s1.area_per_channel[236] + s1.area_per_channel[237] + s1.area_per_channel[243]
+
         event_data['s1_area_upper_injection_fraction'] = area_upper_injection / s1.area
         event_data['s1_area_lower_injection_fraction'] = area_lower_injection / s1.area
-        
+
         # Want S1 AreaFractionTop Probability
         if s1.area < self.low_pe_threshold:
             s1_frac = s1.area / self.low_pe_threshold
@@ -180,7 +180,8 @@ def extract_data(self, event):
                 statistic=self.statistic)
 
             pmts_bottom = np.setdiff1d(self.tpc_channels, confused_s1_channels)
-            pmts_bottom[0:127] = 0
+            pmts_bottom[self.pax_config['DEFAULT']['channels_top']] = 0
+
             event_data['s1_pattern_fit_bottom_hax'] = self.pattern_fitter.compute_gof(
                 (self.x[i], self.y[i], self.z[i]),
                 apc[self.tpc_channels],