SPE spec bug fix, adding trigger truth

wfsim/core/pulse.py

@@ -29,7 +29,11 @@ def __init__(self, config):
         self.init_pmt_current_templates()
         self.init_spe_scaling_factor_distributions()
         self.config['turned_off_pmts'] = np.arange(len(config['gains']))[np.array(config['gains']) == 0]
-
+        self.current_max = np.max(np.array(self._pmt_current_templates), axis=1)
+        self.current_2_adc = self.config['pmt_circuit_load_resistor'] \
+            * self.config['external_amplification'] \
+            / (self.config['digitizer_voltage_range'] / 2 ** (self.config['digitizer_bits']))
+
         self.clear_pulse_cache()
 
     def __call__(self, *args):
@@ -52,7 +56,11 @@ def __call__(self, *args):
                                                      len(self._photon_timings)).astype(np.int64)
 
         dt = self.config.get('sample_duration', 10)  # Getting dt from the lib just once
-        self._n_double_pe = self._n_double_pe_bot = 0  # For truth aft output
+        self._n_photon = self._n_photon_bottom = 0  # For truth output
+        self._n_pe = self._n_pe_bottom = 0
+        self._n_photon_trigger = self._n_photon_trigger_bottom = 0
+        self._n_pe_trigger = self._n_pe_trigger_bottom = 0
+        self._raw_area = self._raw_area_bottom = 0
 
         counts_start = 0  # Secondary loop index for assigning channel
         for channel, counts in zip(*np.unique(self._photon_channels, return_counts=True)):
@@ -69,24 +77,27 @@ def __call__(self, *args):
             if '_photon_gains' not in self.__dict__:
 
                 _channel_photon_gains = self.config['gains'][channel] \
-                    * self.uniform_to_pe_arr(np.random.random(len(_channel_photon_timings)))
+                    * self.uniform_to_pe_arr(np.random.random(len(_channel_photon_timings)), channel)
 
                 # Add some double photoelectron emission by adding another sampled gain
                 n_double_pe = np.random.binomial(len(_channel_photon_timings),
                                                  p=self.config['p_double_pe_emision'])
-                self._n_double_pe += n_double_pe
-                if channel in self.config['channels_bottom']:
-                    self._n_double_pe_bot += n_double_pe
-
-                if self.config['detector'] == 'XENON1T':
-                    _channel_photon_gains[:n_double_pe] += self.config['gains'][channel] \
-                        * self.uniform_to_pe_arr(np.random.random(n_double_pe), channel)
-                else:
-                    _channel_photon_gains[:n_double_pe] += self.config['gains'][channel] \
-                        * self.uniform_to_pe_arr(np.random.random(n_double_pe))
+
+                _channel_photon_gains[:n_double_pe] += self.config['gains'][channel] \
+                    * self.uniform_to_pe_arr(np.random.random(n_double_pe), channel)
+
             else:
+                n_double_pe = 0
                 _channel_photon_gains = np.array(self._photon_gains[self._photon_channels == channel])
 
+            # += truth per channel to internal truth fields
+            self.add_truth(
+                _channel_photon_timings,
+                _channel_photon_gains,
+                dt,
+                channel,
+                n_double_pe,)
+
             # Build a simulated waveform, length depends on min and max of photon timings
             min_timing, max_timing = np.min(
                 _channel_photon_timings), np.max(_channel_photon_timings)
@@ -199,6 +210,41 @@ def uniform_to_pe_arr(self, p, channel=0):
         indices = (p * 2000).astype(np.int64) + 1
         return self.__uniform_to_pe_arr[channel, indices]
 
+    def add_truth(self,
+                  photon_timings,
+                  photon_gains,
+                  dt,
+                  channel,
+                  n_double_pe,
+                 ):
+
+        if channel in self.config['turned_off_pmts']:
+            return
+
+        if str(channel) in self.config.get('special_thresholds', {}):
+            threshold = self.config['special_thresholds'][str(channel)] - 0.5
+        else:
+            threshold = self.config['zle_threshold'] - 0.5
+
+        reminder = (photon_timings % dt).astype(int)
+        above_threshold = photon_gains * self.current_max[reminder] * self.current_2_adc > threshold
+        trigger_photon = np.sum(above_threshold)
+        trigger_dpe = np.sum(above_threshold[:n_double_pe])
+        raw_area = np.sum(photon_gains[above_threshold]) / self.config['gains'][channel]
+
+        self._n_photon += len(photon_timings)
+        self._n_pe += len(photon_timings) + n_double_pe
+        self._n_photon_trigger += trigger_photon
+        self._n_pe_trigger += trigger_photon + trigger_dpe
+        self._raw_area += raw_area
+
+        if channel in self.config['channels_bottom']:
+            self._n_photon_bottom += len(photon_timings)
+            self._n_pe_bottom += len(photon_timings) + n_double_pe
+            self._n_photon_trigger_bottom += trigger_photon
+            self._n_pe_trigger_bottom += trigger_photon + trigger_dpe
+            self._raw_area_bottom += raw_area
+
     def clear_pulse_cache(self):
         self._pulses = []
 
@@ -246,6 +292,17 @@ def add_current(photon_timings,
         pulse_current[start:start + template_length] += \
             pmt_current_templates[reminder] * gain_total
 
+    @staticmethod
+    def hit_above_threshold(photon_timings,
+                            photon_gains,
+                            dt,
+                            pmt_current_templates,
+                            config,):
+
+        reminder = int(photon_timings % dt)
+
+
+
     @staticmethod
     def singlet_triplet_delays(size, singlet_ratio, config, phase):
         """

wfsim/core/rawdata.py

@@ -342,19 +342,10 @@ def get_truth(self, instruction, truth_buffer):
                 (self.config['samples_before_pulse_center'] + self.config['samples_after_pulse_center'] + 1) \
                 * self.config['sample_duration']
 
-        channels = getattr(pulse, '_photon_channels', [])
-
-        n_dpe = getattr(pulse, '_n_double_pe', 0)
-        n_dpe_bot = getattr(pulse, '_n_double_pe_bot', 0)
-
-        tb['n_photon'] -= np.sum(np.isin(channels, getattr(pulse, 'turned_off_pmts', [])))
-        tb['n_pe'] += tb['n_photon']+n_dpe
-        # this turned_off guy, check how this works with a config['turned_off_guys']
-        channels_bottom = list(
-            set(self.config['channels_bottom']).difference(getattr(pulse, 'turned_off_pmts', [])))
-        tb['n_photon_bottom'] = (
-            np.sum(np.isin(channels, channels_bottom)))
-        tb['n_pe_bottom'] = tb['n_photon_bottom'] + n_dpe_bot
+        # Copy single valued fields directly from pulse class
+        for field in ['n_pe', 'n_pe_trigger', 'n_photon', 'n_photon_trigger', 'raw_area']:
+            for suffix in ['', '_bottom']:
+                tb[field+suffix] = getattr(pulse, '_' + field + suffix, 0)
 
         # Summarize the instruction cluster in one row of the truth file
         for field in instruction.dtype.names:

wfsim/strax_interface.py

@@ -43,11 +43,17 @@
 
 truth_extra_dtype = [
     (('End time of the interaction [ns]', 'endtime'), np.int64),
-    (('Number of simulated electrons', 'n_electron'), np.float64),
-    (('Number of detected photons', 'n_photon'), np.float64),
-    (('Number of detected photoelectrons(n_photons+dpe)', 'n_pe'), np.float64),
-    (('number of photons detected in bottom array', 'n_photon_bottom'), np.float64),
-    (('number of photoelectrons detected in bottom array', 'n_pe_bottom'), np.float64),
+    (('Number of simulated electrons', 'n_electron'), np.int32),
+    (('Number of photons reaching PMT', 'n_photon'), np.int32),
+    (('Number of photons + dpe passing', 'n_pe'), np.int32),
+    (('Number of photons passing trigger', 'n_photon_trigger'), np.int32),
+    (('Number of photons + dpe passing trigger', 'n_pe_trigger'), np.int32),
+    (('Raw area in pe passing trigger', 'raw_area'), np.float64),
+    (('Number of photons reaching PMT (bottom)', 'n_photon_bottom'), np.int32),
+    (('Number of photons + dpe passing (bottom)', 'n_pe_bottom'), np.int32),
+    (('Number of photons passing trigger (bottom)', 'n_photon_trigger_bottom'), np.int32),
+    (('Number of photons + dpe passing trigger (bottom)', 'n_pe_trigger_bottom'), np.int32),
+    (('Raw area in pe passing trigger (bottom)', 'raw_area_bottom'), np.float64),
     (('Arrival time of the first photon [ns]', 't_first_photon'), np.float64),
     (('Arrival time of the last photon [ns]', 't_last_photon'), np.float64),
     (('Mean time of the photons [ns]', 't_mean_photon'), np.float64),