ZH: Rewritten generate_beacon.py

simulations/airshower_beacon_simulation/generate_beacon.py

@@ -3,6 +3,9 @@
 import numpy as np
 import json
 import h5py
+import os.path as path
+
+from copy import deepcopy as copy
 
 from earsim import REvent, Antenna
 
@@ -10,9 +13,9 @@ import lib
 
 tx_fname = 'tx.json'
 clocks_fname = 'clocks.csv'
-beacon_fname = 'beacons.hdf5'
+antennas_fname = 'antennas.hdf5'
 
-def write_beacon_file(fname, tx, f_beacon):
+def write_tx_file(fname, tx, f_beacon):
     with open(fname, 'w') as fp:
         return json.dump(
                 dict(
@@ -27,7 +30,7 @@ def write_beacon_file(fname, tx, f_beacon):
                 fp
                 )
 
-def read_beacon_file(fname):
+def read_tx_file(fname):
     with open(fname, 'r') as fp:
         data = json.load(fp)
 
@@ -43,6 +46,9 @@ def read_beacon_hdf5(fname):
 
         mydict = { k:tx_attrs.get(k) for k in ['x', 'y', 'z', 'name'] }
         tx = Antenna(**mydict)
+
+        print(f_beacon, tx)
+
         antennas = []
         for k, ant in h5['antennas'].items():
             mydict = { k:ant.attrs.get(k) for k in ['x', 'y', 'z', 'name'] }
@@ -54,19 +60,24 @@ def read_beacon_hdf5(fname):
 
     return f_beacon, tx, antennas
 
-def init_beacon_hdf5(fname, tx, f_beacon):
+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:
             tx_group = fp.create_group('tx')
-        tx_group['f_beacon'] = f_beacon
-        tx_group['x'] = tx.x
-        tx_group['y'] = tx.y
-        tx_group['z'] = tx.z
-        tx_group['name'] = tx.name
+            tx_attrs = tx_group.attrs
+
+            if f_beacon is not None:
+                tx_attrs['f_beacon'] = f_beacon
+
+            if tx is not None:
+                tx_attrs['x'] = tx.x
+                tx_attrs['y'] = tx.y
+                tx_attrs['z'] = tx.z
+                tx_attrs['name'] = tx.name
 
     return fname
 
-def append_beacon_hdf5(fname, antenna, *datasets, clock_offset=None, overwrite=False):
-
+def append_antenna_hdf5(fname, antenna, datasets = [], overwrite=False):
     with h5py.File(fname, 'a') as fp:
         if 'antennas' not in fp.keys():
             group = fp.create_group('antennas')
@@ -77,12 +88,11 @@ def append_beacon_hdf5(fname, antenna, *datasets, clock_offset=None, overwrite=F
             raise NotImplementedError
 
         ant_group = group.create_group(antenna.name)
-        ant_group['x'] = antenna.x
-        ant_group['y'] = antenna.y
-        ant_group['z'] = antenna.z
-        ant_group['name'] = antenna.name
-        if clock_offset is not None:
-            ant_group['clock_offset'] = clock_offset
+        ant_attrs = ant_group.attrs
+        ant_attrs['x'] = antenna.x
+        ant_attrs['y'] = antenna.y
+        ant_attrs['z'] = antenna.z
+        ant_attrs['name'] = antenna.name
 
         dset = ant_group.create_dataset('traces', (len(datasets)+1, len(antenna.t)), dtype='f')
 
@@ -90,16 +100,50 @@ def append_beacon_hdf5(fname, antenna, *datasets, clock_offset=None, overwrite=F
         for i, mydset in enumerate(datasets,1):
             dset[i] = mydset
 
+def BeaconedREvent(fname, fname_beacon=None, beacon_amplitude=(1,0,0)):
+
+    if fname_beacon is None:
+        fname_beacon = path.join(path.dirname(fname), beacon_fname)
+
+    f_beacon, tx, antennas = read_beacon_hdf5(fname_beacon)
+
+    ev = REvent(fname)
+    ev.tx = tx
+    ev.tx.f_beacon = f_beacon
+
+    merge_beacon_into_REvent(antennas, ev, beacon_amplitude=beacon_amplitude)
+    
+    return ev
+    
+def merge_beacon_into_REvent(beacons, ev, beacon_amplitude=(1,0,0)):
+    assert len(beacons) == len(ev.antennas)
+
+    if not hasattr(beacon_amplitude, '__len__'):
+        beacon_amplitude = np.repeat(beacon_amplitude,3)
+
+    assert len(beacon_amplitude) == 3
+
+    for i, beacon in enumerate(beacons):
+        ant = ev.antennas[i]
+        assert len(beacon.t) == len(ant.t)
+
+        # copy original Efields
+        ev.antennas[i].orig_Ex = copy(ev.antennas[i].Ex)
+        ev.antennas[i].orig_Ey = copy(ev.antennas[i].Ey)
+        ev.antennas[i].orig_Ez = copy(ev.antennas[i].Ez)
+
+        ev.antennas[i].beacon = beacon.beacon
+        if beacon_amplitude[0] != 0:
+            ev.antennas[i].Ex += beacon_amplitude[0]*beacon.beacon
+        if beacon_amplitude[1] != 0:
+            ev.antennas[i].Ey += beacon_amplitude[1]*beacon.beacon
+        if beacon_amplitude[2] != 0:
+            ev.antennas[i].Ez += beacon_amplitude[2]*beacon.beacon
 
 if __name__ == "__main__":
     from os import path
 
-    max_clock_offset = 100# ns
     remake_tx = False
-    remake_clock_offsets = False
-
-    seed = 12345
-    rng = np.random.default_rng(seed)
 
     fname = "ZH_airshower/mysim.sry"
     tx = Antenna(x=-500,y=0,z=0,name='tx')
@@ -108,31 +152,29 @@ if __name__ == "__main__":
     ####
     fname_dir = path.dirname(fname)
     tx_fname = path.join(fname_dir, tx_fname)
-    clocks_fname = path.join(fname_dir, clocks_fname)
-    beacon_fname = path.join(fname_dir, beacon_fname)
+    antennas_fname = path.join(fname_dir, antennas_fname)
 
     if not path.isfile(tx_fname) or remake_tx:
-        write_beacon_file(tx_fname, tx, f_beacon)
+        write_tx_file(tx_fname, tx, f_beacon)
     else:
-        tx, f_beacon = read_beacon_file(tx_fname)
-
+        tx, f_beacon = read_tx_file(tx_fname)
 
     # read in antennas
     ev = REvent(fname)
     N_antennas = len(ev.antennas)
 
-    if not path.isfile(clocks_fname) or remake_clock_offsets:
-        clock_offsets = max_clock_offset * (2*rng.uniform(size=N_antennas) - 1)
-        np.savetxt(clocks_fname, clock_offsets)
-    else:
-        clock_offsets = np.loadtxt(clocks_fname)
+    # initialize hdf5 file
+    init_antenna_hdf5(antennas_fname, tx, f_beacon)
 
     # make beacon per antenna
-    init_beacon_hdf5(beacon_fname, tx, f_beacon)
     for i, antenna in enumerate(ev.antennas):
-        append_beacon_hdf5(
-                beacon_fname,
-                antenna,
-                lib.beacon_from(tx, antenna, f_beacon, antenna.t, t0=clock_offsets[i]),
-                clock_offset=clock_offsets[i],
-            )
+        beacon = lib.beacon_from(tx, antenna, f_beacon, antenna.t)
+
+        E = [antenna.Ex, antenna.Ey, antenna.Ez ]
+
+        # add to relevant polarisation
+        E[0] += beacon
+
+        append_antenna_hdf5( antennas_fname, antenna, (E[0], E[1], E[2] ))
+
+    print("Antenna HDF5 file written as " + str(antennas_fname))