photoelectron gain mean redefinition

FlamTeam · Jan 25, 2024 · 30a94ad · 30a94ad
1 parent 9cf24a4
commit 30a94ad
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 15 deletions.
diff --git a/flamedisx/lxe_blocks/raw_signals.py b/flamedisx/lxe_blocks/raw_signals.py
@@ -27,15 +27,21 @@ class MakeFinalSignals(fd.Block):
     signal_name: str
 
     def _simulate(self, d):
+        if self.quanta_name == 'photoelectron':
+            gain= self.gimme_numpy(self.quanta_name + '_gain_mean')
+        else:
+            gain = self.gimme_numpy(self.quanta_name + '_gain_mean', d[self.quanta_name + 's_detected'].values)                
         d[self.signal_name] = stats.norm.rvs(
             loc=(d[self.quanta_name + 's_detected']
-                * self.gimme_numpy(self.quanta_name + '_gain_mean',d[self.quanta_name + 's_detected'].values)
-                ),
+                * gain),
             scale=(d[self.quanta_name + 's_detected']**0.5
                 * self.gimme_numpy(self.quanta_name + '_gain_std')))
 
     def _annotate(self, d):
-        m = self.gimme_numpy(self.quanta_name + '_gain_mean', d['s2_raw_mle'].values)
+        if self.quanta_name == 'photoelectron':
+            m = self.gimme_numpy(self.quanta_name + '_gain_mean')
+        else:
+            m = self.gimme_numpy(self.quanta_name + '_gain_mean',d['s2_raw_mle'].values/20.)
         s = self.gimme_numpy(self.quanta_name + '_gain_std')
 
         mle = d[self.quanta_name + 's_detected_mle'] = \
@@ -55,10 +61,15 @@ def _compute(self,
                  quanta_detected, s_observed,
                  data_tensor, ptensor):
         # Lookup signal gain mean and std per detected quanta
-        mean_per_q = self.gimme(self.quanta_name + '_gain_mean',
-                                bonus_arg=quanta_detected,
+        if self.quanta_name == 'photoelectron':
+            mean_per_q = self.gimme(self.quanta_name + '_gain_mean',
                                 data_tensor=data_tensor,
                                 ptensor=ptensor)[:, o, o]
+        else:
+            mean_per_q = self.gimme(self.quanta_name + '_gain_mean',
+                               bonus_arg=quanta_detected,
+                               data_tensor=data_tensor,
+                               ptensor=ptensor)
         std_per_q = self.gimme(self.quanta_name + '_gain_std',
                                data_tensor=data_tensor,
                                ptensor=ptensor)[:, o, o]
@@ -87,15 +98,14 @@ class MakeS1(MakeFinalSignals):
     signal_name = 's1_raw'
 
     dimensions = ('photoelectrons_detected', 's1_raw')
-    special_model_functions = ('photoelectron_gain_mean',)
-    model_functions = ('photoelectron_gain_std',) + special_model_functions
+    special_model_functions = ()
+    model_functions = (
+        'photoelectron_gain_mean',
+        'photoelectron_gain_std',) + special_model_functions
 
     max_dim_size = {'photoelectrons_detected': 120}
 
-    @staticmethod
-    def photoelectron_gain_mean(phd):
-        return tf.ones_like(phd)
-
+    photoelectron_gain_mean = 1.
     photoelectron_gain_std = 0.5
 
     def _compute(self, data_tensor, ptensor,
@@ -120,7 +130,7 @@ class MakeS2(MakeFinalSignals):
 
     @staticmethod
     def electron_gain_mean(eld):
-        return tf.ones_like(eld)
+        return 20. + 0. * eld
 
     electron_gain_std = 5.
 
@@ -129,4 +139,4 @@ def _compute(self, data_tensor, ptensor,
         return super()._compute(
             quanta_detected=electrons_detected,
             s_observed=s2_raw,
-            data_tensor=data_tensor, ptensor=ptensor)
+            data_tensor=data_tensor, ptensor=ptensor)
diff --git a/flamedisx/xenon/x1t_sr1.py b/flamedisx/xenon/x1t_sr1.py
@@ -500,12 +500,12 @@ def electron_detection_eff(drift_time,
         return extraction_eff * tf.exp(-drift_time / elife)
 
     @staticmethod
-    def electron_gain_mean(s2_relative_ly,
+    def electron_gain_mean(electrons_detected, s2_relative_ly,
                            se_gain_relative,
                            *,
                            single_electron_gain=DEFAULT_SINGLE_ELECTRON_GAIN):
 
-        return single_electron_gain*s2_relative_ly*se_gain_relative
+        return single_electron_gain*s2_relative_ly*se_gain_relative + 0.*electrons_detected
 
     @staticmethod
     def electron_gain_std(s2_relative_ly,