m-thesis-introduction/simulations/airshower_beacon_simulation/generate_beacon.py

179 lines
5 KiB
Python
Executable file

#!/usr/bin/env python3
import numpy as np
import json
import h5py
import os.path as path
from copy import deepcopy as copy
from earsim import REvent, Antenna
import lib
tx_fname = 'tx.json'
antennas_fname = 'antennas.hdf5'
def write_tx_file(fname, tx, f_beacon):
with open(fname, 'w') as fp:
return json.dump(
dict(
f_beacon=f_beacon,
tx=dict(
x=tx.x,
y=tx.y,
z=tx.z,
name=tx.name
)
),
fp
)
def read_tx_file(fname):
with open(fname, 'r') as fp:
data = json.load(fp)
f_beacon = data['f_beacon']
tx = Antenna(**data['tx'])
return tx, f_beacon
def read_beacon_hdf5(fname):
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)
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'] }
antenna = Antenna(**mydict)
antenna.t = ant['traces'][0]
antenna.beacon = ant['traces'][1]
antennas.append(antenna)
return f_beacon, tx, antennas
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_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_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')
else:
group = fp['antennas']
if overwrite:
raise NotImplementedError
ant_group = group.create_group(antenna.name)
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')
dset[0] = antenna.t
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
remake_tx = False
fname = "ZH_airshower/mysim.sry"
tx = Antenna(x=-500,y=0,z=0,name='tx')
f_beacon = 50e-3 # GHz
####
fname_dir = path.dirname(fname)
tx_fname = path.join(fname_dir, tx_fname)
antennas_fname = path.join(fname_dir, antennas_fname)
if not path.isfile(tx_fname) or remake_tx:
write_tx_file(tx_fname, tx, f_beacon)
else:
tx, f_beacon = read_tx_file(tx_fname)
# read in antennas
ev = REvent(fname)
N_antennas = len(ev.antennas)
# initialize hdf5 file
init_antenna_hdf5(antennas_fname, tx, f_beacon)
# make beacon per antenna
for i, antenna in enumerate(ev.antennas):
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))