clipping probability value

FlamTeam · Jan 10, 2024 · 6c3d990 · 6c3d990
1 parent 2a24bd1
commit 6c3d990
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flamedisx/lxe_blocks/s2_loss.py b/flamedisx/lxe_blocks/s2_loss.py
@@ -26,15 +26,15 @@ def _compute(self, data_tensor, ptensor,
         # s2_raw_after_loss distributed as Binom(s2_raw, p=s2_survival_probability)
         result = tfp.distributions.Binomial(
                 total_count=s2_raw,
-                probs=tf.cast(s2_survival_probability, dtype=fd.float_type())
+                probs=tf.clip_by_value(tf.cast(s2_survival_probability, dtype=fd.float_type()), 0.,1.)
             ).prob(s2_raw_after_loss)
 
         return result
 
     def _simulate(self, d):
         d['s2_raw_after_loss'] = stats.binom.rvs(
             n=d['s2_raw'],
-            p=self.gimme_numpy('s2_survival_p'))
+            p=np.nan_to_num(self.gimme_numpy('s2_survival_p')).clip(0., 1.))
 
     def _annotate(self, d):
         # TODO: this assumes the spread from the double PE effect is subdominant