fixing bug

FlamTeam · Jan 10, 2024 · 01a264f · 01a264f
1 parent c8a37d2
commit 01a264f
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)
         s2_raw_after_loss = tf.clip_by_value(s2_raw_after_loss, 1e-15, tf.float32.max)
         result = tfp.distributions.Binomial(
-                total_count=tf.cast(s2_raw, dtype=fd.int_type()),
+                total_count=tf.clip_by_value(tf.cast(s2_raw, dtype=fd.int_type()),1e-15, tf.int32.max),
                 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'].astype(dtype=np.int32),
+            n=np.clip(d['s2_raw'].astype(dtype=np.int32),0,np.iinfo(np.int32).max),
             p=np.nan_to_num(self.gimme_numpy('s2_survival_p')).clip(0., 1.))
 
     def _annotate(self, d):