Fix S1/S2 correction and gas gain when simulating S1/S2PE

appletree/maps/_gas_gain_relative.json

@@ -0,0 +1 @@
+{"coordinate_lowers": [-70, -70], "coordinate_uppers": [70, 70], "coordinate_type": "regbin", "coordinate_name": ["x", "y"], "map": [[1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0]]}

appletree/plugins/detector.py

@@ -13,52 +13,50 @@
 @export
 @takes_config(
     Map(
-        name="s1_correction",
+        name="s1_lce",
         default="_s1_correction.json",
-        help="S1 light collection efficiency correction",
+        help="S1 light collection efficiency",
     ),
 )
-class S1Correction(Plugin):
-    depends_on = ["rec_x", "rec_y", "rec_z"]
-    provides = ["s1_correction"]
+class S1LCE(Plugin):
+    depends_on = ["x", "y", "z"]
+    provides = ["s1_lce"]
 
     @partial(jit, static_argnums=(0,))
-    def simulate(self, key, parameters, rec_x, rec_y, rec_z):
-        pos = jnp.stack([rec_x, rec_y, rec_z]).T
-        s1_correction = self.s1_correction.apply(pos)
-        return key, s1_correction
+    def simulate(self, key, parameters, x, y, z):
+        pos_true = jnp.stack([x, y, z]).T
+        s1_lce = self.s1_lce.apply(pos_true)
+        return key, s1_lce
 
 
 @export
 @takes_config(
     Map(
-        name="s2_correction",
+        name="s2_lce",
         default="_s2_correction.json",
-        help="S2 light collection efficiency correction",
+        help="S2 light collection efficiency",
     ),
 )
-class S2Correction(Plugin):
-    depends_on = ["rec_x", "rec_y"]
-    provides = ["s2_correction"]
+class S2LCE(Plugin):
+    depends_on = ["x", "y"]
+    provides = ["s2_lce"]
 
     @partial(jit, static_argnums=(0,))
-    def simulate(self, key, parameters, rec_x, rec_y):
-        pos = jnp.stack([rec_x, rec_y]).T
-        s2_correction = self.s2_correction.apply(pos)
-        return key, s2_correction
+    def simulate(self, key, parameters, x, y):
+        pos_true = jnp.stack([x, y]).T
+        s2_lce = self.s2_lce.apply(pos_true)
+        return key, s2_lce
 
 
 @export
 class PhotonDetection(Plugin):
-    depends_on = ["num_photon", "s1_correction"]
+    depends_on = ["num_photon", "s1_lce"]
     provides = ["num_s1_phd"]
     parameters = ("g1", "p_dpe")
 
     @partial(jit, static_argnums=(0,))
-    def simulate(self, key, parameters, num_photon, s1_correction):
-        g1_true_no_dpe = jnp.clip(
-            parameters["g1"] * s1_correction / (1.0 + parameters["p_dpe"]), 0, 1.0
-        )
+    def simulate(self, key, parameters, num_photon, s1_lce):
+        g1_true_no_dpe = jnp.clip(parameters["g1"] * s1_lce / (1.0 + parameters["p_dpe"]), 0, 1.0)
         key, num_s1_phd = randgen.binomial(key, g1_true_no_dpe, num_photon)
         return key, num_s1_phd
 
@@ -104,21 +102,28 @@ def simulate(self, key, parameters, num_electron, drift_survive_prob):
         return key, num_electron_drifted
 
 
+@takes_config(
+    Map(
+        name="gas_gain_relative",
+        default="_gas_gain_relative.json",
+        help="Gas gain as a function of (x, y) / mean gas gain",
+    ),
+)
 @export
 class S2PE(Plugin):
-    depends_on = ["num_electron_drifted", "s2_correction"]
+    depends_on = ["num_electron_drifted", "s2_lce", "x", "y"]
     provides = ["num_s2_pe"]
     parameters = ("g2", "gas_gain")
 
     @partial(jit, static_argnums=(0,))
-    def simulate(self, key, parameters, num_electron_drifted, s2_correction):
-        extraction_eff = parameters["g2"] / parameters["gas_gain"]
-        g2_true = parameters["g2"] * s2_correction
-        gas_gain_true = g2_true / extraction_eff
+    def simulate(self, key, parameters, num_electron_drifted, s2_lce, x, y):
+        pos_true = jnp.stack([x, y]).T
+        gas_gain = parameters["gas_gain"] * self.gas_gain_relative.apply((pos_true))
+        extraction_eff = parameters["g2"] * s2_lce / gas_gain
 
         key, num_electron_extracted = randgen.binomial(key, extraction_eff, num_electron_drifted)
 
-        mean_s2_pe = num_electron_extracted * gas_gain_true
+        mean_s2_pe = num_electron_extracted * gas_gain
         key, num_s2_pe = randgen.truncate_normal(key, mean_s2_pe, jnp.sqrt(mean_s2_pe), vmin=0)
 
         return key, num_s2_pe

appletree/plugins/reconstruction.py

@@ -69,6 +69,44 @@ def simulate(self, key, parameters, num_s2_pe):
         return key, s2_area
 
 
+@export
+@takes_config(
+    Map(
+        name="s1_correction",
+        default="_s1_correction.json",
+        help="S1 xyz correction on reconstructed positions",
+    ),
+)
+class S1Correction(Plugin):
+    depends_on = ["rec_x", "rec_y", "rec_z"]
+    provides = ["s1_correction"]
+
+    @partial(jit, static_argnums=(0,))
+    def simulate(self, key, parameters, rec_x, rec_y, rec_z):
+        pos_rec = jnp.stack([rec_x, rec_y, rec_z]).T
+        s1_correction = self.s1_correction.apply(pos_rec)
+        return key, s1_correction
+
+
+@export
+@takes_config(
+    Map(
+        name="s2_correction",
+        default="_s2_correction.json",
+        help="S2 xy correction on constructed positions",
+    ),
+)
+class S2Correction(Plugin):
+    depends_on = ["rec_x", "rec_y"]
+    provides = ["s2_correction"]
+
+    @partial(jit, static_argnums=(0,))
+    def simulate(self, key, parameters, rec_x, rec_y):
+        pos_rec = jnp.stack([rec_x, rec_y]).T
+        s2_correction = self.s2_correction.apply(pos_rec)
+        return key, s2_correction
+
+
 @export
 class cS1(Plugin):
     depends_on = ["s1_area", "s1_correction"]