From e57403e765926d7a93782e0d2cd788756f5c15e3 Mon Sep 17 00:00:00 2001
From: Eric Teunis de Boone <ericteunis@deboone.nl>
Date: Tue, 22 Nov 2022 11:06:15 +0100
Subject: [PATCH] ZH: rewrite read_beacon_hdf5 function

---
 .../aa_generate_beacon.py                     | 59 ++++++++++---------
 1 file changed, 32 insertions(+), 27 deletions(-)

diff --git a/simulations/airshower_beacon_simulation/aa_generate_beacon.py b/simulations/airshower_beacon_simulation/aa_generate_beacon.py
index 3820bb7..2dae7eb 100755
--- a/simulations/airshower_beacon_simulation/aa_generate_beacon.py
+++ b/simulations/airshower_beacon_simulation/aa_generate_beacon.py
@@ -42,41 +42,46 @@ def read_tx_file(fname):
 
     return tx, f_beacon
 
-def read_beacon_hdf5(fname, traces_key='traces', raise_exception=True, read_AxB=True):
+def read_beacon_hdf5(fname, **h5ant_kwargs):
     with h5py.File(fname, 'r') as h5:
-        tx_attrs = h5['tx'].attrs
-        f_beacon = tx_attrs.get('f_beacon')
-
-        mydict = { k:tx_attrs.get(k) for k in ['x', 'y', 'z', 'name'] }
-        tx = Antenna(**mydict)
+        tx = Antenna_from_h5ant(h5['tx'], traces_key=None)
+        f_beacon = tx.attrs['f_beacon']
 
         antennas = []
         for k, h5ant in h5['antennas'].items():
-            mydict = { k:h5ant.attrs.get(k) for k in ['x', 'y', 'z', 'name'] }
-            ant = Antenna(**mydict)
-
-            if traces_key not in h5ant:
-                if raise_exception:
-                    raise ValueError("Traces_key not in file")
-            else:
-                ant.t = h5ant[traces_key][0]
-                ant.Ex = h5ant[traces_key][1]
-                ant.Ey = h5ant[traces_key][2]
-                ant.Ez = h5ant[traces_key][3]
-                if len(h5ant[traces_key]) > 4:
-                    ant.beacon = h5ant[traces_key][4]
-
-                if read_AxB and 'E_AxB' in h5ant:
-                    ant.t_AxB = h5ant['E_AxB'][0]
-                    ant.E_AxB = h5ant['E_AxB'][1]
-
-            if h5ant.attrs:
-                ant.attrs = {**h5ant.attrs}
+            ant = Antenna_from_h5ant(h5ant, **h5ant_kwargs)
 
             antennas.append(ant)
 
     return f_beacon, tx, antennas
 
+def Antenna_from_h5ant(h5ant, traces_key='traces', raise_exception=True, read_AxB=True):
+    mydict = { k:h5ant.attrs.get(k) for k in ['x', 'y', 'z', 'name'] }
+    ant = Antenna(**mydict)
+
+    if traces_key is None:
+        pass
+    elif traces_key not in h5ant:
+        if raise_exception:
+            raise ValueError("Traces_key not in file")
+    else:
+        ant.t = h5ant[traces_key][0]
+        ant.Ex = h5ant[traces_key][1]
+        ant.Ey = h5ant[traces_key][2]
+        ant.Ez = h5ant[traces_key][3]
+        if len(h5ant[traces_key]) > 4:
+            ant.beacon = h5ant[traces_key][4]
+
+        if read_AxB and 'E_AxB' in h5ant:
+            ant.t_AxB = h5ant['E_AxB'][0]
+            ant.E_AxB = h5ant['E_AxB'][1]
+
+    if h5ant.attrs:
+        ant.attrs = {**h5ant.attrs}
+
+    return ant
+
+
 def init_antenna_hdf5(fname, tx = None, f_beacon = None):
     with h5py.File(fname, 'w') as fp:
         if tx is not None or f_beacon is not None:
@@ -135,7 +140,7 @@ def append_antenna_hdf5(fname, antenna, columns = [], name='traces', prepend_tim
 
         for i, col in enumerate(columns, 1*prepend_time):
             dset[i] = col
-    
+
 if __name__ == "__main__":
     from os import path