expand plotting.py

plotting.py

@@ -5,14 +5,21 @@ import numpy as np
 import matplotlib.pyplot as plt
 from matplotlib.collections import PatchCollection
 from matplotlib.colors import Normalize
-from utils import distance2showeraxis
-from pylab import *
 import seaborn.apionly as sns
 
+import utils
 
-def plot_array(v_stations, values, label='', vmin=None, vmax=None):
-    """Plot a map *values* for an detector array specified by *v_stations*.
-    """
+
+def maybe_save(fig, fname):
+    if fname is not None:
+        fig.savefig(fname, bbox_inches='tight')
+        plt.close()
+
+
+def plot_array(v_stations, values, v_core=None, v_axis=None,
+              label='', title='', vmin=None, vmax=None, fname=None):
+    """Plot a map *values* for an detector array specified by *v_stations*. """
+    print('Plot event')
     xd, yd, zd = v_stations.T / 1000  # in [km]
     circles = [plt.Circle((x, y), 0.650) for x, y in zip(xd.flat, yd.flat)]
     fig, ax = plt.subplots(1)
@@ -21,30 +28,36 @@ def plot_array(v_stations, values, label='', vmin=None, vmax=None):
     coll.set_array(values)
     coll.cmap.set_under('#d3d3d3')
     ax.add_collection(coll)
-    ax.set(aspect='equal')
-    ax.grid()
     cbar = fig.colorbar(coll)
     cbar.set_label(label)
+    ax.grid(True)
+    ax.set_title(title)
+    ax.set_aspect('equal')
     ax.set_xlim(-8, 8)
     ax.set_ylim(-8, 8)
     ax.set_xlabel('x [km]')
     ax.set_ylabel('y [km]')
+    # plot shower direction and core
+    if (v_core is not None) and (v_axis is not None):
+        x, y, z = v_core / 1000 - 3 * v_axis
+        dx, dy, dz = 3 * v_axis
+        plt.arrow(x, y, dx, dy, lw=2, head_width=0.4, head_length=0.5, fc='r', ec='r')
+    maybe_save(fig, fname)
 
 
-def plot_traces_of_array_for_one_event(Smu, Sem, v_axis, v_stations, n=5):
-    """ Plot time traces of the n^2 central tanks
-        Sem = EM-SigmalTrace, Smu = Muon-SignalTrace
-    """
-    n0 = int(len(v_stations)**.5)
-    i0 = (n0 - n) // 2
+def plot_array_traces(Smu, Sem, v_stations, n=5, fname=None):
+    """ Plot time traces of the n^2 central tanks. """
+    print('Plot time traces event')
+    n0 = int(len(v_stations)**.5)  # number of stations along one axis
+    i0 = (n0 - n) // 2  # start index for the sub-array to be plotted
 
     S1 = Smu.reshape(n0, n0, -1)
     S2 = Sem.reshape(n0, n0, -1)
+    v = v_stations.reshape(n0, n0, 3)
 
-    fig, axes = subplots(n, n, sharex=True, figsize=(29, 16), facecolor='w')
+    fig, axes = plt.subplots(n, n, sharex=True, figsize=(29, 16), facecolor='w')
     plt.tight_layout()
     t = np.arange(12.5, 2001, 25)
-    v_stations = v_stations.reshape(11,11,3)
     for ix in range(n):
         for iy in range(n):
             ax = axes[ix, iy]
@@ -53,57 +66,123 @@ def plot_traces_of_array_for_one_event(Smu, Sem, v_axis, v_stations, n=5):
             ax.step(t, h1 + h2, c='k', where='mid')
             ax.step(t, h1, label='$\mu$', where='mid')
             ax.step(t, h2, label='$e\gamma$', where='mid')
-            title = "X:"+ str(v_stations[ix + i0,iy + i0][0])+ " / Y:"+ str(v_stations[ix + i0,iy + i0][1])            
-            ax.legend(title=title + " / r=" + str(ix + i0) + ", " + str(iy + i0), fontsize='x-small')
-            ax.set_xlim(0, 1500)
+            ax.legend(fontsize='x-small')
             ax.grid(True)
+            ax.set_title('%.1f, %.1f km' % tuple(v[ix + i0, iy + i0, 0:2] / 1000))
             ax.set_xlabel('$t$ / ns', fontsize='x-small')
-    ax.set_ylabel('$S$ / VEM', fontsize='x-small')
-    savefig('trace_VEM.png')
+            ax.set_ylabel('$S$ / VEM', fontsize='x-small')
+            ax.set_xlim(0, 1500)
+    maybe_save(fig, fname)
 
-def plot_time_distribution(t):
-    # histogram of time values
-    print "time"
+
+def plot_time_distribution(T, fname=None):
+    """ histogram of time values """
+    print('Plot time distribution')
     fig, ax = plt.subplots(1)
-    ax.hist(t[~np.isnan(t)].flatten(), bins=40, normed=True)
+    ax.hist(T[~np.isnan(T)].flatten() * 1E6, bins=40, normed=True)
     ax.grid()
-    ax.set_xlabel('time [ms]')
-    ax.set_ylabel('frequency')
-    fig.savefig('./time_distribution.png')
+    ax.set_xlabel('time [$\mu$ s]')
+    ax.set_ylabel('relative frequency')
+    maybe_save(fig, fname)
+
 
-def plot_total_signal_distribution(s):
-    print "signal"
-    # histogram of signal values
-    s = np.sum(s, axis=-1)
+def plot_signal_distribution(S, fname=None):
+    """ histogram of signal values """
+    print('Plot total signal distribution')
+    s = np.sum(S, axis=-1)  # sum over time trace per station
     s[np.isnan(s)] = 0
+    s = np.sum(s, axis=-1)  # sum over stations
     fig, ax = plt.subplots(1)
-    ax.hist(np.sum(s, axis=(-1)), bins=401, normed=True)
+    ax.hist(np.log10(s), bins=31, normed=True)
     ax.grid()
-    plt.yscale('log')
-    ax.set_xlabel('signal [a.u.]')
-    ax.set_ylabel('frequency')
-    fig.savefig('./total_signal_distribution.png')
+    ax.set_xlabel('$\log_{10}$(total signal) [a.u.]')
+    ax.set_ylabel('relative frequency')
+    maybe_save(fig, fname)
+
 
-def plot_energy_distribution(logE):
-    print "energy"
-    # histogram of Energy values
+def plot_energy_distribution(logE, fname=None):
+    """ histogram of energy values """
+    print('Plot energy distribution')
     fig, ax = plt.subplots(1)
-    ax.hist(logE, bins = np.linspace(18.5, 20, 69), normed=False)
+    ax.hist(logE, bins=np.linspace(18.5, 20, 31))
     ax.grid()
     ax.set_xlabel('energy [eV]')
     ax.set_ylabel('frequency')
-    fig.savefig("./energy_distribution.png")
+    maybe_save(fig, fname)
+
+
+def plot_zenith_distribution(zenith, fname=None):
+    """ histogram of zenith values """
+    print('Plot zenith distribution')
+    fig, ax = plt.subplots(1)
+    ax.hist(np.rad2deg(zenith), bins=np.linspace(0, 60, 31))
+    ax.grid()
+    ax.set_xlabel('zenith [degree]')
+    ax.set_ylabel('frequency')
+    maybe_save(fig, fname)
+
+
+def plot_phi_distribution(phi, fname=None):
+    """ histogram of phi values """
+    print('Plot phi distribution')
+    fig, ax = plt.subplots(1)
+    ax.hist(np.rad2deg(phi), bins=np.linspace(-180, 180, 31))
+    ax.set_xticks([-180, -90, 0, 90, 180])
+    ax.grid()
+    ax.set_xlabel('phi [degree]')
+    ax.set_ylabel('frequency')
+    maybe_save(fig, fname)
 
-def plot_energy_vs_tanks_with_hits(logE, s):
-    print "energy vc tanks"
-    # histogram tons with hits vs. Energy
-    nt = np.sum(~np.isnan(s), axis=(1))
+
+def plot_xmax_distribution(Xmax, fname=None):
+    """ histogram of Xmax values """
+    print('Plot Xmax distribution')
+    fig, ax = plt.subplots(1)
+    ax.hist(Xmax, bins=np.linspace(600, 1200, 31))
+    ax.grid()
+    ax.set_xlabel('Xmax [g/cm$^2$]')
+    ax.set_ylabel('frequency')
+    maybe_save(fig, fname)
+
+
+def plot_stations_vs_energy(logE, S, fname=None):
+    """ histogram of stations with hits vs energy """
+    print('Plot stations vs energy')
+    nt = np.sum(~np.isnan(S), axis=1)
     fig = plt.figure()
-    ax = sns.regplot(x=logE, y=nt,  x_bins = np.linspace(18.5, 20, 69), fit_reg=False)
+    bins = np.linspace(18.5, 20, 31)
+    ax = sns.regplot(x=logE, y=nt, x_bins=bins, fit_reg=False)
     ax.grid()
+    ax.set_ylim(0)
     ax.set_xlabel('Energy')
     ax.set_ylabel('Number of Stations')
-    fig.savefig('./energy_vs_nstations.png')
+    maybe_save(fig, fname)
+
+
+def plot_stations_vs_zenith(zen, S, fname=None):
+    """ histogram of stations with hits vs zenith """
+    print('Plot stations vs zenith')
+    nt = np.sum(~np.isnan(S), axis=1)
+    fig = plt.figure()
+    ax = sns.regplot(x=np.rad2deg(zen), y=nt, x_bins=np.linspace(0, 60, 31), fit_reg=False)
+    ax.grid()
+    ax.set_ylim(0)
+    ax.set_xlabel('Zenith')
+    ax.set_ylabel('Number of Stations')
+    maybe_save(fig, fname)
+
+
+def plot_stations_vs_phi(phi, S, fname=None):
+    """ histogram of stations with hits vs zenith """
+    print('Plot stations vs phi')
+    nt = np.sum(~np.isnan(S), axis=1)
+    fig = plt.figure()
+    ax = sns.regplot(x=np.rad2deg(phi), y=nt, x_bins=np.linspace(-180, 180, 31), fit_reg=False)
+    ax.grid()
+    ax.set_ylim(0)
+    ax.set_xlabel('Phi')
+    ax.set_ylabel('Number of Stations')
+    maybe_save(fig, fname)
 
 
 if __name__ == '__main__':
@@ -117,10 +196,40 @@ if __name__ == '__main__':
     v_stations = d['detector']
     T = d['time']
     S = d['signal']
+    S1 = d['signal1']
+    S2 = d['signal2']
+    phi, zenith = utils.vec2ang(v_axis)
 
+    # ------------------------------------
     # plot example event
-    plot_array(v_stations, T[0] * 1E6, label='time [mu s]', vmin=-8, vmax=8)
-    plt.savefig('time.png')
-    logS = np.log10(S.sum(axis=-1))
-    plot_array(v_stations, logS[0], label='log(signal)', vmin=0, vmax=5)
-    plt.savefig('signal.png')
+    # ------------------------------------
+    for i in range(3):
+        title = 'logE=%.2f, Xmax=%.2f, zenith=%.2f' % (logE[i], Xmax[i], np.rad2deg(zenith[i]))
+
+        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,
+            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='time [mu s]', title=title,
+            fname='plots/example-%i-signal.png' % i)
+
+        plot_array_traces(
+            Smu=S1[i], Sem=S2[i], v_stations=v_stations, n=5,
+            fname='plots/example-%i-traces.png' % i)
+
+    # ------------------------------------
+    # plot distribution of all events
+    # ------------------------------------
+    plot_time_distribution(T, fname='plots/time_distribution.png')
+    plot_signal_distribution(S, fname='plots/signal_distribution.png')
+    plot_energy_distribution(logE, fname='plots/energy_distribution.png')
+    plot_xmax_distribution(Xmax, fname='plots/xmax_distribution.png')
+    plot_zenith_distribution(zenith, fname='plots/zenith_distribution.png')
+    plot_phi_distribution(phi, fname='plots/phi_distribution.png')
+    plot_stations_vs_energy(logE, S, fname='plots/stations_vs_energy.png')
+    plot_stations_vs_zenith(zenith, S, fname='plots/stations_vs_zenith.png')
+    plot_stations_vs_zenith(phi, S, fname='plots/stations_vs_phi.png')

shower.py

@@ -7,7 +7,6 @@ import gumbel
 import plotting
 
 
-atm = atmosphere.Atmosphere()
 HEIGHT = 1400
 SPACING = 1500
 
@@ -24,7 +23,7 @@ def station_response(r, dX, logE=19, A=1):
     """
     # strength of muon and em signal
     # scaling with energy
-    S0 = 2000 * 10**(logE - 19)
+    S0 = 1200 * 10**(logE - 19)
     # scaling with mass number and relative scaling mu/em
     a = A**0.15
     S1 = S0 * a / (a + 1) * 1
@@ -126,7 +125,7 @@ if __name__ == '__main__':
 
     # showers
     print('simulating showers')
-    nb_events = 100
+    nb_events = 1000
     logE = 18.5 + 1.5 * np.random.rand(nb_events)
     # logE = 20 * np.ones(nb_events)
     mass = 1 * np.ones(nb_events)
@@ -139,49 +138,43 @@ 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, s2 = detector_response(logE[i], mass[i], v_axis[i], v_core[i], v_max[i], v_stations)
-        S1[i] = s1
-        S2[i] = s2
+        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])
 
+    # 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
-#    S1 += 0.02 * S1 * np.random.randn(*S1.shape)
-#    S2 += 0.02 * S2 * np.random.randn(*S2.shape)
-#    # add absolute noise to signal pulse
-#    S1 += 0.5 + 0.2 * np.random.randn(*S1.shape)
-#    S2 += 0.5 + 0.2 * np.random.randn(*S2.shape)
+    # add relative noise to signal pulse
+    S += 0.02 * S * np.random.randn(*S.shape)
+
+    # add absolute noise to signal pulse
+    S += 1 + 0.2 * np.random.randn(*S.shape)
 
     # trigger threshold: use only stations with sufficient signal-to-noise
-##    c = S.sum(axis=-1) < 80 * 0.55
-##    T[c] = np.NaN
-##    S[c] = np.NaN
+    c = S.sum(axis=-1) < 80 * 1.2
+    T[c] = np.NaN
+    S[c] = np.NaN
 
     # TODO: apply array trigger (3T5)
 
     # save
-#    np.savez_compressed(
-#        'showers.npz',
-#        logE=logE,
-#        mass=mass,
-#        Xmax=Xmax,
-#        time=T,
-#        signal=S,
-#        showercore=v_core,
-#        showeraxis=v_axis,
-#        showermax=v_max,
-#        detector=v_stations)
-
-    plotting.plot_traces_of_array_for_one_event(
-        Smu=S1[0], Sem=S2[0], v_axis=v_axis[0], v_stations=v_stations, n=5)
-
-    # r = utils.distance2showeraxis(v_stations, v_core[0], v_axis[0])
-    # plotting.plot_array(v_stations, r)
-
-    import matplotlib.pyplot as plt
-    plt.show()
+    np.savez_compressed(
+        'showers.npz',
+        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)