diff --git a/simulations/airshower_beacon_simulation/ba_measure_beacon_phase.py b/simulations/airshower_beacon_simulation/ba_measure_beacon_phase.py
index 0037897..a98bcb6 100755
--- a/simulations/airshower_beacon_simulation/ba_measure_beacon_phase.py
+++ b/simulations/airshower_beacon_simulation/ba_measure_beacon_phase.py
@@ -97,11 +97,15 @@ if __name__ == "__main__":
                 # TODO: refine masking
                 # use beacon but remove where E_AxB-Beacon != 0
                 if True:
-                    if not True:
-                        t_mask = np.isclose(h5ant['E_AxB'][1], h5ant['traces'][4], rtol=1e-3, atol=1e-3)
-                    else:
-                        t_mask = np.ones(len(t_trace), dtype=bool)
-                        t_mask[1500:3000] = False # magic numbers from aa_generate_beacon
+                    N_pre, N_post = 250, 250 # TODO: make this configurable
+
+                    max_idx = np.argmax(test_traces[0])
+
+                    low_idx = max(0, max_idx-N_pre)
+                    high_idx = min(len(t_trace), max_idx+N_post)
+
+                    t_mask = np.ones(len(t_trace), dtype=bool)
+                    t_mask[low_idx:high_idx] = False
 
                     t_trace = t_trace[t_mask]
                     for j, t in enumerate(test_traces):
diff --git a/simulations/airshower_beacon_simulation/ca_period_from_shower.py b/simulations/airshower_beacon_simulation/ca_period_from_shower.py
index a76a733..cfa76ee 100755
--- a/simulations/airshower_beacon_simulation/ca_period_from_shower.py
+++ b/simulations/airshower_beacon_simulation/ca_period_from_shower.py
@@ -213,7 +213,7 @@ if __name__ == "__main__":
         ev.antennas[i].t_AxB += measured_repair_offsets[i]
 
         if apply_signal_window_from_max:
-            N_pre, N_post = 250, 250
+            N_pre, N_post = 250, 250 # TODO: make this configurable
 
             max_idx = np.argmax(ant.E_AxB)
 
diff --git a/simulations/airshower_beacon_simulation/dc_grid_power_time_fixes.py b/simulations/airshower_beacon_simulation/dc_grid_power_time_fixes.py
index 54035b4..8bb4cce 100755
--- a/simulations/airshower_beacon_simulation/dc_grid_power_time_fixes.py
+++ b/simulations/airshower_beacon_simulation/dc_grid_power_time_fixes.py
@@ -133,7 +133,7 @@ if __name__ == "__main__":
 
     # Slice the traces to a small part around the peak
     if apply_signal_window_from_max:
-        N_pre, N_post = 250, 250
+        N_pre, N_post = 250, 250 # TODO: make this configurable
 
         for i, ant in enumerate(ev.antennas):
             max_idx = np.argmax(ant.E_AxB)