adding block for electron loss at wall

flamedisx/lxe_blocks/detection.py

@@ -59,10 +59,15 @@ def _simulate(self, d):
                 'penning_quenching_eff', d['photons_produced'].values)
         else:
             p *= self.gimme_numpy(
-                'electron_loss', d['electrons_produced'].values)
+                'electron_loss', d['electrons_survived'].values)
+
+        if self.quanta_name == 'photon':
+            n = d[self.quanta_name + 's_produced']
+        else:
+            n = d['electrons_survived']
 
         d[self.quanta_name + 's_detected'] = stats.binom.rvs(
-            n=d[self.quanta_name + 's_produced'],
+            n=n,
             p=p)
         d['p_accepted'] *= self.gimme_numpy(
             self.quanta_name + '_acceptance',
@@ -85,8 +90,12 @@ def _annotate(self, d):
                              "configure your cuts correctly?")
 
         # Estimate produced quanta
-        n_prod_mle = d[self.quanta_name + 's_produced_mle'] = \
-            d[self.quanta_name + 's_detected_mle'] / eff
+        if self.quanta_name == 'photon':
+            n_prod_mle = d[self.quanta_name + 's_produced_mle'] = \
+                d[self.quanta_name + 's_detected_mle'] / eff
+        else:
+            n_prod_mle = d['electrons_survived_mle'] = \
+                d[self.quanta_name + 's_detected_mle'] / eff
 
         # Estimating the spread in number of produced quanta is tricky since
         # the number of detected quanta is itself uncertain.
@@ -96,9 +105,14 @@ def _annotate(self, d):
 
         for bound, sign, intify in (('min', -1, np.floor),
                                     ('max', +1, np.ceil)):
-            d[self.quanta_name + 's_produced_' + bound] = intify(
-                n_prod_mle + sign * self.source.max_sigma * _std
-            ).clip(0, None).astype(int)
+            if self.quanta_name == 'photon':
+                d[self.quanta_name + 's_produced_' + bound] = intify(
+                    n_prod_mle + sign * self.source.max_sigma * _std
+                ).clip(0, None).astype(int)
+            else:
+                d['electrons_survived_' + bound] = intify(
+                    n_prod_mle + sign * self.source.max_sigma * _std
+                ).clip(0, None).astype(int)
 
 
 @export
@@ -133,12 +147,12 @@ def _compute(self, data_tensor, ptensor,
 
 @export
 class DetectElectrons(DetectPhotonsOrElectrons):
-    dimensions = ('electrons_produced', 'electrons_detected')
+    dimensions = ('electrons_survived', 'electrons_detected')
 
     special_model_functions = ('electron_acceptance', 'electron_loss')
     model_functions = ('electron_detection_eff',) + special_model_functions
 
-    max_dim_size = {'electrons_produced': 100}
+    max_dim_size = {'electrons_survived_mle': 100}
 
     @staticmethod
     def electron_detection_eff(drift_time, *,
@@ -157,7 +171,7 @@ def electron_loss(nel):
     quanta_name = 'electron'
 
     def _compute(self, data_tensor, ptensor,
-                 electrons_produced, electrons_detected):
-        return super()._compute(quanta_produced=electrons_produced,
+                 electrons_survived, electrons_detected):
+        return super()._compute(quanta_produced=electrons_survived,
                                 quanta_detected=electrons_detected,
                                 data_tensor=data_tensor, ptensor=ptensor)

flamedisx/lxe_blocks/electron_wall_loss.py

@@ -0,0 +1,70 @@
+import typing as ty
+
+import numpy as np
+from scipy import stats
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+import flamedisx as fd
+export, __all__ = fd.exporter()
+o = tf.newaxis
+
+@export
+class ElectronWallLoss(fd.Block):
+    dimensions = ('electrons_produced', 'electrons_survived')
+
+    #special_model_functions = ('electron_survival_probability')
+    model_functions = ('electron_survival_probability',) #+ special_model_functions
+
+    max_dim_size = {'electrons_produced': 100}
+
+    quanta_name = 'electron'
+
+    electron_survival_probability = 1.
+
+
+    def _compute(self, data_tensor, ptensor,
+                 electrons_produced, electrons_survived):
+        p = self.gimme('electron_survival_probability',
+                       data_tensor=data_tensor, ptensor=ptensor)[:, o, o]
+
+        result = tfp.distributions.Binomial(
+                total_count=quanta_produced,
+                probs=tf.cast(p, dtype=fd.float_type())
+            ).prob(electrons_survived)
+
+        return result
+
+    def _simulate(self, d):
+        p = self.gimme_numpy('electron_survival_probability')
+
+        d['electrons_survived'] = stats.binom.rvs(
+            n=d['electrons_produced'], p=p)
+
+
+    def _annotate(self, d):
+        # Get efficiency
+        eff = self.gimme_numpy('electron_survival_probability')
+
+        # Check for bad efficiencies
+        if self.check_efficiencies and np.any(eff <= 0):
+            raise ValueError(f"Found event with nonpositive "
+                             "electron_survival_probability: did you apply and "
+                             "configure your inputs correctly?")
+
+        # Estimate produced electrons
+        n_prod_mle = d[self.quanta_name + 's_produced_mle'] = \
+            d['electrons_survived'] / eff
+
+        # Estimating the spread in number of produced quanta is tricky since
+        # the number of produced_after_loss quanta is itself uncertain.
+        # TODO: where did this derivation come from again?
+        q = (1 - eff) / eff
+        _std = (q + (q ** 2 + 4 * n_prod_mle * q) ** 0.5) / 2
+
+        for bound, sign, intify in (('min', -1, np.floor),
+                                    ('max', +1, np.ceil)):
+            d[self.quanta_name + 's_produced_' + bound] = intify(
+                n_prod_mle + sign * self.source.max_sigma * _std
+            ).clip(0, None).astype(int)
+