add spherical wavefront

plotting.py

@@ -208,13 +208,13 @@ if __name__ == '__main__':
 
         plot_array(
             v_stations, T[i] * 1E6, v_core=v_core[i], v_axis=v_axis[i],
-            vmin=-10, vmax=10, label='time [mu s]', title=title,
+            vmin=-10, vmax=10, label='time [$\mu$ s]', title=title,
             fname='plots/example-%i-time.png' % i)
 
         logStot = np.log10(S.sum(axis=-1))
         plot_array(
             v_stations, logStot[i], v_core=v_core[i], v_axis=v_axis[i],
-            vmin=1, vmax=5, label='log(Total_Signal)', title=title,
+            vmin=1, vmax=5, label='$\log_{10}(S_\mathrm{tot})$', title=title,
             fname='plots/example-%i-signal.png' % i)
 
         plot_array_traces(

shower.py

@@ -22,8 +22,7 @@ def station_response(r, dX, logE=19, zenith=0, mass=1):
     Returns:
         time traces f_muon(t) and f_em(t) in [VEM] for t = 0 - 2000 ns in bins of 25 ns
     """
-    # signal strength
-    # scaling with energy
+    # signal strength, scaling with energy
     S0 = 900 * 10**(logE - 19)
     # # scaling with zenith angle (similar to S1000 --> S38 relation, CIC)
     # x = np.cos(zenith)**2 - np.cos(np.deg2rad(38))**2
@@ -36,8 +35,8 @@ def station_response(r, dX, logE=19, zenith=0, mass=1):
     S1 *= (np.maximum(r, 150) / 1000)**-4.7
     S2 *= (np.maximum(r, 150) / 1000)**-6.1
     # scaling with traversed atmosphere to station
-    S1 *= np.minimum((dX / 100)**-0.1, 10) ## -0.1
-    S2 *= np.minimum((dX / 100)**-0.4, 10) ## -0.4
+    S1 *= np.minimum((dX / 100)**-0.1, 10)
+    S2 *= np.minimum((dX / 100)**-0.4, 10)
 
     # limit total signal, otherwise we get memory problems when drawing that many samples from distribution
     Stot = S1 + S2
@@ -67,8 +66,7 @@ def station_response(r, dX, logE=19, zenith=0, mass=1):
 
 
 def detector_response(logE, mass, v_axis, v_core, v_max, v_stations):
-    """ Simulate the detector response for all SD stations and one event
-    """
+    """ Simulate the detector response for all SD stations and one event. """
     _, zenith = utils.vec2ang(v_axis)
 
     r = utils.distance2showeraxis(v_stations, v_core, v_axis)  # radial distance to shower axis
@@ -90,50 +88,49 @@ def detector_response(logE, mass, v_axis, v_core, v_max, v_stations):
     return S1, S2
 
 
-def rand_shower_geometry(N, logE, mass):
-    """ Generate random shower geometries: Xmax, axis, core, maximum. """
-    N = len(logE)
+def rand_shower_geometry(logE, mass):
+    """ Generate random shower geometries: Xmax, axis, core, maximum, virtual origin """
+    nb_showers = len(logE)
 
     # 1) random shower axis
-    phi = 2 * np.pi * (np.random.rand(N) - 0.5)
-    zenith = utils.rand_zenith(N)
+    phi = 2 * np.pi * (np.random.rand(nb_showers) - 0.5)
+    zenith = utils.rand_zenith(nb_showers)
     v_axis = utils.ang2vec(phi, zenith)
 
     # 2) random shower core on ground (offset w.r.t. grid origin)
-    v_core = SPACING * (np.random.rand(N, 3) - 0.5)
+    v_core = SPACING * (np.random.rand(nb_showers, 3) - 0.5)
     v_core[:, 2] = HEIGHT  # core is always at ground level
 
-    # 3) random shower maximum, require h > hmin
-    Xmax = np.empty(N)
-    hmax = np.empty(N)
-    i = 0
-    j = 0
-    while i < N:
-        j += 1
-        Xmax[i] = gumbel.rand_gumbel(logE[i], mass[i])
-        hmax[i] = atmosphere.height_at_slant_depth(Xmax[i], zenith[i])
-        if hmax[i] > HEIGHT + 200:
-            i += 1
-    print('%i of %i showers accepted' % (i, j))
-    dmax = (hmax - HEIGHT) / np.cos(zenith)  # distance to Xmax
-    v_max = v_core + v_axis * dmax[:, np.newaxis]
+    # 3) random shower maximum, require Xmax 200m above ground
+    Xmax = gumbel.rand_gumbel(logE, mass)
+    Xmax_max = atmosphere.slant_depth(HEIGHT + 200, zenith)
 
-    return Xmax, v_axis, v_core, v_max
+    while not np.all(Xmax < Xmax_max):
+        idx = Xmax > Xmax_max  # resample these Xmax values
+        Xmax[idx] = gumbel.rand_gumbel(logE[idx], mass[idx])
 
+    # 4) point of shower maximum
+    h = atmosphere.height_at_slant_depth(Xmax, zenith)
+    d = (h - HEIGHT) / np.cos(zenith)
+    v_max = v_core + v_axis * d[:, np.newaxis]
 
-if __name__ == '__main__':
-    # detector array
-    n = 11
-    v_stations = utils.station_coordinates(n, layout='offset')  # x,y,z coordinates of SD stations
-    nb_stations = n**2  # number of stations
+    # 5) virtual origin at 90% of Xmax
+    h = atmosphere.height_at_slant_depth(0.9 * Xmax, zenith)
+    d = (h - HEIGHT) / np.cos(zenith)
+    v_origin = v_core + v_axis * d[:, np.newaxis]
+
+    return Xmax, v_axis, v_core, v_max, v_origin
+
+
+def rand_events(logE, mass, v_stations, fname=None):
+    """Simulate events for given energy, mass, and detector positions."""
 
-    # showers
+    nb_events = len(logE)
+    nb_stations = len(v_stations)
+
+    # simulation showers
     print('simulating showers')
-    nb_events = 100000
-    logE = 18.5 + 1.5 * np.random.rand(nb_events)
-    # logE = 20 * np.ones(nb_events)
-    mass = 1 * np.ones(nb_events)
-    Xmax, v_axis, v_core, v_max = rand_shower_geometry(nb_events, logE, mass)
+    Xmax, v_axis, v_core, v_max, v_origin = rand_shower_geometry(logE, mass)
 
     # detector response for each event
     print('simulating detector response')
@@ -142,51 +139,59 @@ if __name__ == '__main__':
     S2 = np.zeros((nb_events, nb_stations, 80))
     for i in range(nb_events):
         print('%4i, logE = %.2f' % (i, logE[i]))
-        S1[i], S2[i] = detector_response(
-            logE[i], mass[i], v_axis[i], v_core[i], v_max[i], v_stations)
-        T[i] = utils.arrival_time_planar(v_stations, v_core[i], v_axis[i])
+        S1[i], S2[i] = detector_response(logE[i], mass[i], v_axis[i], v_core[i], v_max[i], v_stations)
+        # T[i] = utils.arrival_time_planar(v_stations, v_core[i], v_axis[i])
+        T[i] = utils.arrival_time_spherical(v_stations, v_origin[i], v_core[i])
 
     # total signal
     S = S1 + S2
-
     # add per ton noise on arrival time
     T += 20E-9 * np.random.randn(*T.shape)
-
     # add time offset per event (to account for core position)
     T += 100E-9 * (np.random.rand(nb_events, 1) - 0.5)
-
     # add relative noise to signal pulse
     S += 0.02 * S * np.random.randn(*S.shape)
-
     # add absolute noise to signal pulse
     noise_level = 1.2
     S += noise_level * (1 + 0.5 * np.random.randn(*S.shape))
-
     # trigger threshold: use only stations with sufficient signal-to-noise
     c = S.sum(axis=-1) < 80 * noise_level * 1.2
     T[c] = np.NaN
     S[c] = np.NaN
 
-    # TODO: apply array trigger (3T5)
-
-    # save
-    limit = 10000
-    print "Saving data..."
-    n = nb_events//limit
-    for i in range(n):
-        np.savez_compressed('showersTEST_%i.npz' %i,logE=logE[limit*i:limit*(i+1)], mass=mass[limit*i:limit*(i+1)], Xmax=Xmax[limit*i:limit*(i+1)], time=T[limit*i:limit*(i+1)], signal=S[limit*i:limit*(i+1)], signal1=S1[limit*i:limit*(i+1)], signal2=S2[limit*i:limit*(i+1)], showercore=v_core[limit*i:limit*(i+1)], showeraxis=v_axis[limit*i:limit*(i+1)], showermax=v_max[limit*i:limit*(i+1)], detector=v_stations[limit*i:limit*(i+1)])
-
-    np.savez_compressed('showers_%i.npz' %n,
- logE=logE[limit*n:], mass=mass[limit*n:], Xmax=Xmax[limit*n:], time=T[limit*n:], signal=S[limit*n:], signal1=S1[limit*n:], signal2=S2[limit*n:], showercore=v_core[limit*n:], showeraxis=v_axis[limit*n:], showermax=v_max[limit*n:], detector=v_stations[limit*n:])
-    phi, zenith = utils.vec2ang(v_axis)
-    plotting.plot_time_distribution(T, fname='plots/time_distribution.png')
-    plotting.plot_signal_distribution(S, fname='plots/signal_distribution.png')
-    plotting.plot_energy_distribution(logE, fname='plots/energy_distribution.png')
-    plotting.plot_xmax_distribution(Xmax, fname='plots/xmax_distribution.png')
-    plotting.plot_zenith_distribution(zenith, fname='plots/zenith_distribution.png')
-    plotting.plot_phi_distribution(phi, fname='plots/phi_distribution.png')
-    plotting.plot_stations_vs_energy(logE, S, fname='plots/stations_vs_energy.png')
-    plotting.plot_stations_vs_zenith(zenith, S, fname='plots/stations_vs_zenith.png')
-    plotting.plot_stations_vs_zenith(phi, S, fname='plots/stations_vs_phi.png')
-    plotting.plot_array_traces(Smu=S1[0], Sem=S2[0], v_stations=v_stations, n=5, fname='plots/example-trace.png')
-print "Finished!"
+    return {
+        'logE': logE,
+        'mass': mass,
+        'Xmax': Xmax,
+        'time': T,
+        'signal': S,
+        'signal1': S1,
+        'signal2': S2,
+        'showercore': v_core,
+        'showeraxis': v_axis,
+        'showermax': v_max,
+        'detector': v_stations}
+
+
+if __name__ == '__main__':
+    # detector array, vector of (x,y,z) positions
+    v_stations = utils.station_coordinates(11, layout='offset')
+
+    # simulate events
+    nb_events = 100
+    logE = 18.5 + 1.5 * np.random.rand(nb_events)
+    mass = 1 * np.ones(nb_events)
+    data = rand_events(logE, mass, v_stations)
+    np.savez_compressed(
+        'showers.npz',
+        logE=data['logE'],
+        mass=data['mass'],
+        Xmax=data['Xmax'],
+        time=data['time'],
+        signal=data['signal'],
+        signal1=data['signal1'],
+        signal2=data['signal2'],
+        showercore=data['showercore'],
+        showeraxis=data['showeraxis'],
+        showermax=data['showermax'],
+        detector=data['detector'])

utils.py

@@ -110,7 +110,7 @@ def arrival_time_planar(v, vc, va):
     """ Get arrival times for a planar wavefront.
     Note: The shower core is not considered here and as it only adds a constant time offset to all stations.
     Args:
-        v (N x 3 array)
+        v (N x 3 array): array of vectors
         vc (3 array): shower core
         va (3 array): shower axis, pointing upwards
     Return:
@@ -118,3 +118,18 @@ def arrival_time_planar(v, vc, va):
     """
     d = distance2showerplane(v, vc, -va)
     return d / 3E8
+
+
+def arrival_time_spherical(v, v0, vc=None):
+    """ Get arrival times for a spherical wavefront.
+    Args:
+        v (N x 3 array): array of vectors
+        v0 (3 array): virtual shower origin
+        vc (3 array, optional): shower core
+    Return:
+        array: arrival times [s]
+    """
+    d = np.linalg.norm(v - v0, axis=1)
+    if vc is not None:
+        d -= np.linalg.norm(vc - v0)
+    return d / 3E8