fixing shape issue

FlamTeam · Jan 23, 2024 · 2e41c8e · 2e41c8e
1 parent 26f09e2
commit 2e41c8e
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/flamedisx/lxe_blocks/raw_signals.py b/flamedisx/lxe_blocks/raw_signals.py
@@ -127,14 +127,16 @@ class MakeS2(MakeFinalSignals):
 
     max_dim_size = {'electrons_detected': 120}
 
+    @staticmethod
+    def geometrical_acceptance(z, *, geo_acc=1):
+        return geo_acc * tf.ones_like(z)
+
     @staticmethod
     def electron_gain_mean(z, *, g2=20):
         return g2 * tf.ones_like(z)
 
     electron_gain_std = 5.
 
-    geometrical_acceptance = 1.
-
     def _compute(self, data_tensor, ptensor,
                  electrons_detected, s2_raw):
         return super()._compute(