s1 naive z correction

amstrax/plugins/events/corrected_areas.py

@@ -28,6 +28,12 @@ class CorrectedAreas(strax.Plugin):
 
     elife = amstrax.XAMSConfig(default=30000, help="electron lifetime in [ns]")
 
+    s1_naive_z_correction = amstrax.XAMSConfig(default=[-52, 0, 700, -7.69],
+        help="Parameters for the z-dependent S1 correction \
+            [zmin, zmax, y0, a] where y0 + a*z is the correction",
+    )
+
+
     def infer_dtype(self):
         dtype = []
         dtype += strax.time_fields
@@ -43,6 +49,19 @@ def infer_dtype(self):
 
         return dtype
 
+    def s1_naive_z_correction(self, z):
+        """
+        Apply a naive z-dependent S1 correction.
+        Returns the correction factor for the S1 area.
+        """
+
+        s1_correction_function = lambda z: y0 + a * z
+        zmin, zmax, y0, a = self.s1_naive_z_correction
+        s1_correction_average = s1_correction_function((zmin + zmax) / 2)
+        correction = s1_correction_average/s1_correction_function(z)
+
+        return correction
+
     def compute(self, events):
         result = np.zeros(len(events), self.dtype)
         result["time"] = events["time"]
@@ -55,7 +74,7 @@ def compute(self, events):
 
         for peak_type in ["", "alt_"]:
 
-            result[f"{peak_type}cs1"] = events[f"{peak_type}s1_area"]
+            result[f"{peak_type}cs1"] = events[f"{peak_type}s1_area"]*self.s1_naive_z_correction(events["z"])
             result[f"{peak_type}cs2"] = events[f"{peak_type}s2_area"] * np.exp(events["drift_time"] / self.elife)
 
         return result