From 7df39cf6aba8a6713104ef4c299c98e254454fd5 Mon Sep 17 00:00:00 2001
From: Eric Teunis de Boone <ericteunis@deboone.nl>
Date: Fri, 18 Nov 2022 11:02:41 +0100
Subject: [PATCH] ZH: small unassorted tweaks

---
 .../airshower_beacon_simulation/aa_generate_beacon.py  |  7 +++++--
 simulations/airshower_beacon_simulation/lib.py         | 10 ++++++++--
 .../{view_ant0.py => view_orig_ant0.py}                | 10 +++++-----
 3 files changed, 18 insertions(+), 9 deletions(-)
 rename simulations/airshower_beacon_simulation/{view_ant0.py => view_orig_ant0.py} (90%)

diff --git a/simulations/airshower_beacon_simulation/aa_generate_beacon.py b/simulations/airshower_beacon_simulation/aa_generate_beacon.py
index 433bad8..5e7fca9 100755
--- a/simulations/airshower_beacon_simulation/aa_generate_beacon.py
+++ b/simulations/airshower_beacon_simulation/aa_generate_beacon.py
@@ -85,7 +85,7 @@ def init_antenna_hdf5(fname, tx = None, f_beacon = None):
 
     return fname
 
-def append_antenna_hdf5(fname, antenna, columns = [], name='traces', prepend_time=True, overwrite=True):
+def append_antenna_hdf5(fname, antenna, columns = [], name='traces', prepend_time=True, overwrite=True, attrs_dict={}):
     if not overwrite:
         raise NotImplementedError
 
@@ -111,6 +111,9 @@ def append_antenna_hdf5(fname, antenna, columns = [], name='traces', prepend_tim
         ant_attrs['z'] = antenna.z
         ant_attrs['name'] = antenna.name
 
+        for k,v in attrs_dict.items():
+            ant_attrs[k] = v
+
         if name in ant_group:
             if not overwrite:
                 raise NotImplementedError
@@ -141,7 +144,7 @@ if __name__ == "__main__":
         tx = Antenna(x=-20e3,y=0,z=0,name='tx') # m
         dist = lib.distance(tx, Antenna(x=0, y=0, z=0))
 
-        ampl = dist**2
+        ampl = max(1, dist**2)
 
         beacon_amplitudes *= ampl
 
diff --git a/simulations/airshower_beacon_simulation/lib.py b/simulations/airshower_beacon_simulation/lib.py
index b5e1a81..b167219 100644
--- a/simulations/airshower_beacon_simulation/lib.py
+++ b/simulations/airshower_beacon_simulation/lib.py
@@ -1,3 +1,4 @@
+# vim: fdm=indent ts=4
 """
 Library for this simulation
 """
@@ -5,8 +6,11 @@ import numpy as np
 
 from earsim import Antenna
 
-def sine_beacon(f, t, t0=0, amplitude=1, baseline=0):
-    return amplitude * np.sin(2*np.pi*f*(t-t0)) + baseline
+def sine_beacon(f, t, t0=0, amplitude=1, baseline=0, phase=0):
+    """
+    Return a sine appropriate as a beacon
+    """
+    return amplitude * np.cos(2*np.pi*f*(t+t0) + phase) + baseline
 
 
 def distance(x1, x2):
@@ -27,6 +31,8 @@ def beacon_from(tx, rx, f, t=0, t0=0, c_light=3e8, radiate_rsq=True, amplitude=1
     t0 = t0 + dist/c_light
 
     if radiate_rsq:
+        if np.isclose(dist, 0):
+            dist = 1
         amplitude *= 1/(dist**2)
 
     return sine_beacon(f, t, t0=t0, amplitude=amplitude,**kwargs)
diff --git a/simulations/airshower_beacon_simulation/view_ant0.py b/simulations/airshower_beacon_simulation/view_orig_ant0.py
similarity index 90%
rename from simulations/airshower_beacon_simulation/view_ant0.py
rename to simulations/airshower_beacon_simulation/view_orig_ant0.py
index 0ecba30..37ad5a4 100755
--- a/simulations/airshower_beacon_simulation/view_ant0.py
+++ b/simulations/airshower_beacon_simulation/view_orig_ant0.py
@@ -31,7 +31,7 @@ def plot_antenna_Efields(antenna, ax=None, plot_Ex=True, plot_Ey=True, plot_Ez=T
 
     return ax
 
-def plot_antenna_geometry(antennas, ax=None, log_max=False, plot_names=True, plot_max_values=True,**kwargs):
+def plot_antenna_geometry(antennas, ax=None, plot_names=True, plot_max_values=True, log_max=False, colors=None,**kwargs):
     """Show the max values at each antenna."""
     default_kwargs = dict(
             cmap='Spectral_r'
@@ -50,10 +50,10 @@ def plot_antenna_geometry(antennas, ax=None, log_max=False, plot_names=True, plo
 
         if log_max:
             max_vals = np.log10(max_vals)
-    else:
-        max_vals = None
 
-    sc = ax.scatter(x, y, c=max_vals, **kwargs)
+        colors = max_vals
+
+    sc = ax.scatter(x, y, c=colors, **kwargs)
 
     if plot_max_values:
         fig = ax.get_figure()
@@ -67,7 +67,7 @@ def plot_antenna_geometry(antennas, ax=None, log_max=False, plot_names=True, plo
     ax.set_xlabel("[m]")
 
 
-    return ax
+    return ax, sc
 
 if __name__ == "__main__":
     fname = "ZH_airshower/mysim.sry"