trying to fix null differential rates on fake data

FlamTeam · Jan 11, 2024 · 046e640 · 046e640
1 parent 34daa66
commit 046e640
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/flamedisx/lxe_blocks/s2_loss.py b/flamedisx/lxe_blocks/s2_loss.py
@@ -22,19 +22,21 @@ def _compute(self, data_tensor, ptensor,
                  s2_raw, s2_raw_after_loss):
         s2_survival_probability = self.gimme('s2_survival_p',
                            data_tensor=data_tensor, ptensor=ptensor)[:, o, o]
-
+        invalid = s2_raw_after_loss > s2_raw
         # s2_raw_after_loss distributed as Binom(s2_raw, p=s2_survival_probability)
         result = tfp.distributions.Binomial(
-                total_count=tf.cast(s2_raw, dtype=fd.float_type()),
+                total_count=s2_raw,
                 probs=tf.clip_by_value(s2_survival_probability, 0.,1.)
             ).prob(s2_raw_after_loss)
 
-        return result
+        return tf.where(invalid,
+                        tf.zeros_like(s2_raw),
+                        result)
 
     def _simulate(self, d):
         d['s2_raw_after_loss'] = tfp.distributions.Binomial(
-                total_count=tf.cast(d['s2_raw'], dtype=fd.float_type()),
-                probs=tf.clip_by_value(self.gimme_numpy('s2_survival_p'), 0.,1.)
+                total_count=d['s2_raw'],
+                probs=tf.clip_by_value(self.gimme_numpy('s2_survival_p'), 0.,1.)#.astype(np.float64)
             ).sample()
 
     def _annotate(self, d):
@@ -45,5 +47,5 @@ def _annotate(self, d):
             (d['s2_raw_after_loss' + '_mle'] / s2_survival_probability).clip(0, None)
         scale = mle*s2_survival_probability*(1-s2_survival_probability)
 
-        d['s2_raw'  + '_min'] = np.floor(mle-self.source.max_sigma*scale).clip(0, None).astype(int)
-        d['s2_raw'  + '_max'] = np.ceil(mle+self.source.max_sigma*scale).clip(0, None).astype(int)
+        d['s2_raw'  + '_min'] = np.floor(mle-self.source.max_sigma*scale).clip(0, None)
+        d['s2_raw'  + '_max'] = np.ceil(mle+self.source.max_sigma*scale).clip(0, None)
diff --git a/flamedisx/xenon/x1t_sr1.py b/flamedisx/xenon/x1t_sr1.py
@@ -31,7 +31,7 @@
 DEFAULT_G2_TOTAL = DEFAULT_G2 / (1.-DEFAULT_AREA_FRACTION_TOP)
 DEFAULT_SINGLE_ELECTRON_GAIN = DEFAULT_G2_TOTAL / DEFAULT_EXTRACTION_EFFICIENCY
 DEFAULT_SINGLE_ELECTRON_WIDTH = 0.25 * DEFAULT_SINGLE_ELECTRON_GAIN
-DEFAULT_S2_SURVIVAL_PROBABILITY = 1
+DEFAULT_S2_SURVIVAL_PROBABILITY = 1.
 
 # Official numbers from BBF
 DEFAULT_S1_RECONSTRUCTION_BIAS_PIVOT = 0.5948841302444277