diff --git a/simulations/airshower_beacon_simulation/ba_measure_beacon_phase.py b/simulations/airshower_beacon_simulation/ba_measure_beacon_phase.py
index 8b097d4..6cc51f7 100755
--- a/simulations/airshower_beacon_simulation/ba_measure_beacon_phase.py
+++ b/simulations/airshower_beacon_simulation/ba_measure_beacon_phase.py
@@ -23,21 +23,28 @@ if __name__ == "__main__":
 
     from scriptlib import MyArgumentParser
     parser = MyArgumentParser()
-    parser.add_argument('--no-use-AxB', dest='use_AxB_trace', action='store_false')
+
+    group1 = parser.add_mutually_exclusive_group()
+    group1.add_argument('--AxB', dest='use_AxB_trace', action='store_true', help='Only use AxB trace, if both AxB and beacon are not used, we use the antenna polarisations.')
+    group1.add_argument('--beacon', dest='use_beacon_trace', action='store_true', help='Only use the beacon trace')
+
+    parser.add_argument('--N-mask', type=float, default=500, help='Mask N_MASK samples around the absolute maximum of the trace. (Default: %(default)d)')
+
     args = parser.parse_args()
 
     f_beacon_band = (49e-3,55e-3) #GHz
     allow_frequency_fitting = False
     read_frequency_from_file = True
+    N_mask = int(args.N_mask)
 
-    use_AxB_trace = args.use_AxB_trace
-    use_beacon_trace = True # only applicable if AxB = False
+    use_only_AxB_trace = args.use_AxB_trace
+    use_only_beacon_trace = args.use_beacon_trace # only applicable if AxB = False
 
     show_plots = args.show_plots
 
     figsize = (12,8)
 
-    print("use_AxB_trace:", use_AxB_trace, "use_beacon_trace:",use_beacon_trace)
+    print("use_only_AxB_trace:", use_only_AxB_trace, "use_only_beacon_trace:", use_only_beacon_trace)
 
     ####
     fname_dir = args.data_dir
@@ -83,59 +90,64 @@ if __name__ == "__main__":
             h5ant = group[name]
 
             # use E_AxB only instead of polarisations
-            if use_AxB_trace:
-                if 'E_AxB' not in h5ant.keys():
-                    print(f"Antenna does not have 'E_AxB' in {name}")
+            if use_only_AxB_trace:
+                traces_key = 'E_AxB'
+                if traces_key not in h5ant.keys():
+                    print(f"Antenna does not have '{traces_key}' in {name}")
                     sys.exit(1)
 
-                traces = h5ant['E_AxB']
-
+                traces = h5ant[traces_key]
                 t_trace = traces[0]
                 test_traces = [ traces[1] ]
                 orients = ['E_AxB']
 
-                # TODO: refine masking
-                # use beacon but remove where E_AxB-Beacon != 0
-                if True:
-                    N_pre, N_post = 250, 250 # TODO: make this configurable
+            # Only beacon
+            elif use_only_beacon_trace:
+                traces_key = 'filtered_traces'
+                if traces_key not in h5ant.keys():
+                    print(f"Antenna file corrupted? no '{traces_key}' in {name}")
+                    sys.exit(1)
 
-                    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):
-                        test_traces[j] = t[t_mask]
-                        orients[j] = orients[j] + ' masked'
+                traces = h5ant[traces_key]
+                t_trace = traces[0]
+                test_traces = [traces[4]]
+                orients = ['B']
 
             # use separate polarisations
             else:
-                if 'traces' not in h5ant.keys():
-                    print(f"Antenna file corrupted? no 'traces' in {name}")
+                traces_key = 'filtered_traces'
+                if traces_key not in h5ant.keys():
+                    print(f"Antenna file corrupted? no '{traces_key}' in {name}")
                     sys.exit(1)
 
-                traces = h5ant['traces']
+                traces = h5ant[traces_key]
                 t_trace = traces[0]
+                test_traces = [traces[i] for i in range(1,4)]
+                orients = ['Ex', 'Ey', 'Ez']
 
-                if use_beacon_trace:
-                    # only take the Beacon trace
-                    test_traces = [traces[4]]
-                    orients = ['B']
-                else:
-                    test_traces = traces[1:]
-                    orients = ['Ex', 'Ey', 'Ez', 'B']
-
-            # modify the length of the traces
+            # Really only select the first component
             if False:
-                t_trace = t_trace[:len(t_trace)//2]
-                half_traces = []
-                for trace in test_traces:
-                    half_traces.append( trace[:len(trace)//2])
-                test_traces = half_traces
+                test_traces = [test_traces[0]]
+                orients = [orients[0]]
+
+            # TODO: refine masking
+            # use beacon but remove where E_AxB-Beacon != 0
+            # Uses the first traces as reference
+            if N_mask and orients[0] != 'B':
+                N_pre, N_post = N_mask//2, N_mask//2
+
+                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):
+                    test_traces[j] = t[t_mask]
+                    orients[j] = orients[j] + ' masked'
 
             # Do Fourier Transforms
             # to find phases and amplitudes
@@ -155,7 +167,7 @@ if __name__ == "__main__":
                 amps = [ 3e-7 ]
 
             # choose highest amp
-            idx = 0
+            idx = np.argmax(amps)
             if False and len(beacon_phases) > 1:
                 #idx = np.argmax(amplitudes, axis=-1)
                 raise NotImplementedError