WGAN tutorial Shower GAN

.gitignore

+*.pyc

ShowerGAN_signal.py

+import numpy as np
+from tensorflow import keras
+from keras.layers import *
+from keras.layers.merge import _Merge
+from keras.models import Sequential, Model
+from keras.optimizers import Adam
+from keras.layers.advanced_activations import LeakyReLU
+from functools import partial
+from keras.utils import plot_model
+import keras.backend.tensorflow_backend as KTF
+import glob
+import utils
+import tensorflow as tf
+
+
+KTF.set_session(utils.get_session())  # Allows 2 jobs per GPU, Please do not change this during the tutorial
+
+EPOCHS = 10
+GRADIENT_PENALTY_WEIGHT = 10
+BATCH_SIZE = 256
+NCR = 5
+latent_size = 512
+# load trainings data
+Folder = '/home/JGlombitza/DataLink/ToyMC/PreProData' # Training no black tanks
+filenames=glob.glob(Folder + "/PlanarWaveShower_PrePro_2*")
+
+shower_maps, Energy, shower_axis = utils.ReadInData(filenames)
+shower_maps=shower_maps[:,:,:,0]
+np.savez("Data", Energy=Energy, shower_maps=shower_maps)
+shower_maps = shower_maps[:,:,:,1,np.newaxis]
+Energy = Energy/np.max(Energy)
+utils.plot_multiple_signalmaps(shower_maps[:,:,:,0], log_dir=log_dir, title='Footprints', epoch='Real')
+
+N = shower_maps.shape[0]
+
+class RandomWeightedAverage(_Merge):
+    """Takes a randomly-weighted average of two tensors. In geometric terms, this outputs a random point on the line
+    between each pair of input points.
+    Inheriting from _Merge is a little messy but it was the quickest solution I could think of.
+    Improvements appreciated."""
+
+    def _merge_function(self, inputs):
+        weights = K.random_uniform((BATCH_SIZE, 1, 1, 1))
+        return (weights * inputs[0]) + ((1 - weights) * inputs[1])
+
+
+# build generator
+def build_generator(latent_size):
+    inputs = Input(shape=(latent_size,))    
+    x = Dense(latent_size, activation='relu')(inputs)
+    x = Reshape((1,1,latent_size))(x)
+    x = UpSampling2D(size=(3,3))(x)
+    x = Conv2D(64, (2, 2), padding='same', kernel_initializer='he_normal', activation='relu')(x)
+    x = BatchNormalization()(x)
+    x = UpSampling2D(size=(3,3))(x)
+    x = Conv2D(128, (3, 3), padding='same', kernel_initializer='he_normal', activation='relu')(x)
+    x = BatchNormalization()(x)
+    x = Conv2D(128, (3, 3), padding='same', kernel_initializer='he_normal', activation='relu')(x)
+    x = BatchNormalization()(x)
+    x = Conv2D(256, (3, 3), padding='same', kernel_initializer='he_normal', activation='relu')(x)
+    x = BatchNormalization()(x)
+    outputs = Conv2D(1, (3, 3), padding='same', kernel_initializer='he_normal', activation='relu')(x)
+    return Model(inputs=inputs, outputs=outputs, name='generator')
+
+
+def build_critic():
+    critic = Sequential(name='critic')
+    critic.add(Conv2D(64, (3, 3), padding='same', kernel_initializer='he_normal', input_shape=(9,9,1)))
+    critic.add(LeakyReLU())
+    critic.add(Conv2D(128, (3, 3), padding='same', kernel_initializer='he_normal'))
+    critic.add(LeakyReLU())
+    critic.add(Conv2D(128, (3, 3), padding='same', kernel_initializer='he_normal'))
+    critic.add(LeakyReLU())
+    critic.add(Conv2D(256, (3, 3), padding='same', kernel_initializer='he_normal'))
+    critic.add(LeakyReLU())
+    critic.add(GlobalMaxPooling2D())
+    critic.add(Dense(100))
+    critic.add(LeakyReLU())
+    critic.add(Dense(1))
+    return critic
+
+generator = build_generator(latent_size)
+plot_model(generator, to_file=log_dir + '/generator.png', show_shapes=True)
+critic = build_critic()
+plot_model(critic, to_file=log_dir + '/critic.png', show_shapes=True)
+
+# make trainings model for generator
+utils.make_trainable(critic, False)
+utils.make_trainable(generator, True)
+
+generator_in = Input(shape=(latent_size,))
+generator_out = generator(generator_in)
+critic_out = critic(generator_out)
+generator_training = Model(inputs=generator_in, outputs=critic_out)
+generator_training.compile(optimizer=Adam(0.0001, beta_1=0.5, beta_2=0.9, decay=0.0), loss=[utils.wasserstein_loss])
+plot_model(generator_training, to_file=log_dir + '/generator_training.png', show_shapes=True)
+
+# make trainings model for critic
+utils.make_trainable(critic, True)
+utils.make_trainable(generator, False)
+
+generator_in_critic_training = Input(shape=(latent_size,), name="noise")
+shower_in_critic_training = Input(shape=(9,9,1), name='shower_maps')
+generator_out_critic_training = generator(generator_in_critic_training)
+out_critic_training_gen = critic(generator_out_critic_training)
+out_critic_training_shower = critic(shower_in_critic_training)
+averaged_batch = RandomWeightedAverage(name='Average')([generator_out_critic_training, shower_in_critic_training])
+averaged_batch_out = critic(averaged_batch)
+critic_training = Model(inputs=[generator_in_critic_training, shower_in_critic_training], outputs=[out_critic_training_gen, out_critic_training_shower, averaged_batch_out])
+gradient_penalty = partial(utils.gradient_penalty_loss, averaged_batch=averaged_batch, penalty_weight=GRADIENT_PENALTY_WEIGHT)
+gradient_penalty.__name__ = 'gradient_penalty'
+critic_training.compile(optimizer=Adam(0.0001, beta_1=0.5, beta_2=0.9, decay=0.0), loss=[utils.wasserstein_loss, utils.wasserstein_loss, gradient_penalty])
+plot_model(critic_training, to_file=log_dir + '/critic_training.png', show_shapes=True)
+
+# For Wassersteinloss
+positive_y = np.ones(BATCH_SIZE)
+negative_y = -positive_y
+dummy = np.zeros(BATCH_SIZE) # keras throws an error when calculating a loss withuot having a label -> needed for using the gradient penalty loss
+
+generator_loss = []
+critic_loss = []
+
+# trainings loop
+iterations_per_epoch = N//((NCR+1)*BATCH_SIZE)
+for epoch in range(EPOCHS):
+    print "epoch: ", epoch
+    generated_map = generator.predict_on_batch(np.random.randn(BATCH_SIZE, latent_size))
+    utils.plot_multiple_signalmaps(generated_map[:,:,:,0], log_dir=log_dir, title='Generated Footprints Epoch: ', epoch=str(epoch))
+    for iteration in range(iterations_per_epoch):
+        for j in range(NCR):
+            noise_batch = np.random.randn(BATCH_SIZE, latent_size) # generate noise batch for generator
+            shower_batch = shower_maps[BATCH_SIZE*(j+iteration):BATCH_SIZE*(j++iteration+1)]
+            critic_loss.append(critic_training.train_on_batch([noise_batch, shower_batch], [negative_y, positive_y, dummy]))
+            print "critic loss:", critic_loss[-1]
+        noise_batch = np.random.randn(BATCH_SIZE, latent_size) # generate noise batch for generator
+        generator_loss.append(generator_training.train_on_batch([noise_batch], [positive_y]))
+        print "generator loss:", generator_loss[-1]
+
+utils.plot_loss(critic_loss, name="critic", iterations_per_epoch=iterations_per_epoch, NCR=NCR, log_dir=log_dir, bins_per_epoch=3)
+utils.plot_loss(generator_loss, name="generator", iterations_per_epoch=iterations_per_epoch, log_dir=log_dir, bins_per_epoch=3)
+
+
+# plot some generated figures
+generated_map = generator.predict(np.random.randn(BATCH_SIZE, latent_size))
+utils.plot_multiple_signalmaps(generated_map[:,:,:,0], log_dir=log_dir, title='Generated Footprints', epoch='Final')

utils.py

+import numpy as np
+import tensorflow as tf
+from tensorflow import keras
+from keras.layers import *
+from keras.models import Sequential, Model
+from keras.layers.merge import _Merge
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+
+def get_session(gpu_fraction=0.40):
+    ''' Allocate only a fraction of the GPU RAM - (1080 GTX 8Gb)'''
+    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_fraction)
+    return tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
+
+
+def ReadInData(filenames):
+    N = 10000 *len(filenames)
+    a = 10000 # packagesize
+    Input2 = np.zeros(N*9*9*2).reshape(N,9,9,2)
+    Energy = np.zeros(N)
+    showeraxis = np.zeros(N*3).reshape(N,3)
+    
+    for i in range(0, len(filenames)):
+        data = np.load(filenames[i])
+        Energy[a*i:a*(i+1)] = data['Energy']
+        Input2[a*i:a*(i+1)] = data['Input2'].reshape(a,9,9,2)
+        showeraxis[a*i:a*(i+1)] = data['showeraxis']
+    Input2[:,:,:,0] = Input2[:,:,:,0] - np.min(Input2[:,:,:,0], axis=(1,2))[:,np.newaxis,np.newaxis] + 1
+    Filter = Input2[:,:,:,1] == 0
+    Input2[:,:,:,0][Filter] = 0
+    return Input2, Energy, showeraxis
+
+
+def ang2vec(phi, zenith):
+    """ Get 3-vector from spherical angles.
+    Args:
+        phi (array): azimuth (pi, -pi), 0 points in x-direction, pi/2 in y-direction
+        zenith (array): zenith (0, pi), 0 points in z-direction
+    Returns:
+        array of 3-vectors
+    """
+    x = np.sin(zenith) * np.cos(phi)
+    y = np.sin(zenith) * np.sin(phi)
+    z = np.cos(zenith)
+    return np.array([x, y, z]).T
+
+
+def vec2ang(v):
+    '''Calculates phi and theta for given shower axis'''
+    x, y, z = np.asarray(v).T
+    phi = np.arctan2(y, x)
+    theta = np.pi/2 - np.arctan2(z, (x**2 + y**2)**.5)
+    return np.rad2deg(phi), np.rad2deg(theta)
+
+
+def DenselyConnectedConv(z, nfilter):
+    ''' Densely Connected SeparableConvolution2D Layer'''
+    c = Conv2D(nfilter, (3, 3), padding = 'same', activation='relu', kernel_initializer='he_normal')(z)
+    return concatenate([z, c], axis=-1)
+
+
+def make_trainable(model, trainable):
+    ''' Freezes/unfreezes the weights in the given model '''
+    for layer in model.layers:
+        model.trainable=trainable
+    model.trainable=trainable
+
+
+def rectangular_array(n=9):
+     """ Return x,y coordinates for rectangular array with n^2 stations. """
+     n0 = (n - 1) / 2
+     return (np.mgrid[0:n, 0:n].astype(float) - n0)
+
+
+def wasserstein_loss(y_true, y_pred):
+    """Calculates the Wasserstein loss - critic maximises the distance between its output for real and generated samples.
+    To achieve this generated samples have the label -1 and real samples the label 1. Multiplying the outputs by the labels results to the wasserstein loss"""
+    return K.mean(y_true * y_pred)
+
+
+def gradient_penalty_loss(y_true, y_pred, averaged_batch, penalty_weight):
+    """Calculates the gradient penalty (for details see arXiv:1704.00028v3).
+    The 1-Lipschitz constraint for improved WGANs is enforced by adding a term which penalizes if the gradient norm in the critic unequal to 1.
+    (With this loss we penalize gradients with respect to the input of the averaged samples.)"""
+    gradients = K.gradients(K.sum(y_pred), averaged_batch)
+    gradient_l2_norm = K.sqrt(K.sum(K.square(gradients)))
+    gradient_penalty = penalty_weight * K.square(1 - gradient_l2_norm)
+    return gradient_penalty
+
+
+class RandomWeightedAverage(_Merge):
+    """Takes a randomly-weighted average of two tensors"""
+    def __init__(self, batch_size, *args, **kwargs):
+        self.batch_size = batch_size
+        super(_Merge, self).__init__(*args, **kwargs)
+
+    def _merge_function(self, inputs):
+        weights = K.random_uniform((self.batch_size, 1, 1, 1))
+        return (weights * inputs[0]) + ((1 - weights) * inputs[1])
+
+
+def build_generator_graph(generator, critic, conditioner, latent_size, energy_in=1):
+    generator_in = Input(shape=(latent_size,))
+    generator_in_energy_label = Input(shape=(energy_in,))
+    generator_out = generator([generator_in, generator_in_energy_label])
+    critic_out = critic(generator_out)
+    energy_rec_out = conditioner(generator_out)
+    return Model(inputs=[generator_in, generator_in_energy_label], outputs=[critic_out, energy_rec_out])
+
+
+def build_critic_graph(generator, critic, conditioner, latent_size, energy_in, batch_size=1):
+    generator_in_critic_training = Input(shape=(latent_size,), name="noise")
+    generator_in_critic_training_energy_label = Input(shape=(energy_in,), name="energy_label")
+    shower_in_critic_training = Input(shape=(9,9,1), name='shower_maps')
+    generator_out_critic_training = generator([generator_in_critic_training, generator_in_critic_training_energy_label])
+    out_critic_training_gen = critic(generator_out_critic_training)
+    out_critic_training_shower = critic(shower_in_critic_training)
+    averaged_batch = RandomWeightedAverage(batch_size, name='Average')([generator_out_critic_training, shower_in_critic_training])
+    averaged_batch_out = critic(averaged_batch)
+    rec_out = conditioner(shower_in_critic_training)
+    return Model(inputs=[generator_in_critic_training, shower_in_critic_training, generator_in_critic_training_energy_label], outputs=[out_critic_training_gen, out_critic_training_shower, averaged_batch_out, rec_out]), averaged_batch
+
+
+def plot_loss_wasserstein(loss, log_dir = ".", name="", iterations_per_epoch=10, NCR=1, Rec_scaling=1):
+    fig, ax1 = plt.subplots(1, figsize=(10,5))
+    epoch = np.arange(len(loss))
+    loss = np.array(loss)
+    plt.plot(epoch, loss[:,1] + loss[:,2] + loss[:,3], color='red', markersize=12, label=r'$Wasserstein$')
+    plt.legend(loc='upper right',prop={'size':10})
+    ax1.set_xlabel(r'$Iterations$')
+    ax1.set_ylabel(r'$Loss$')
+    ax1.set_ylim(-2, 2)
+    fig.savefig(log_dir + '/%s_Loss_wasserstein.png' %name, dpi=480)
+    plt.plot(epoch, loss[:,4]*Rec_scaling, color='orange', markersize=12, label=r'$Reconstruction$')
+    plt.close('all')
+
+
+def plot_loss_reco(loss, log_dir = ".", name="", iterations_per_epoch=10, NCR=1, Rec_scaling=1):
+    fig, ax1 = plt.subplots(1, figsize=(10,4))
+    epoch = np.arange(len(loss))
+    loss = np.array(loss)
+    plt.plot(epoch, loss[:,2], color='orange', markersize=12, label=r'$Reconstruction$')
+    plt.legend(loc='upper right',prop={'size':10})
+    ax1.set_xlabel(r'$Iterations$')
+    ax1.set_ylabel(r'$Loss$')
+    ax1.set_yscale("log", nonposy='clip')
+    fig.savefig(log_dir + '/%s_Loss_reco.png' %name, dpi=480)
+    plt.close('all')
+
+
+def plot_loss(loss, log_dir = ".", name="", iterations_per_epoch=10, NCR=1, Rec_scaling=1):
+    fig, ax1 = plt.subplots(1, figsize=(10,4))
+    epoch = np.arange(len(loss))
+    loss = np.array(loss)
+    try:
+        plt.plot(epoch, loss[:,0], color='red', markersize=12, label=r'$Total$')
+        plt.plot(epoch, loss[:,1] + loss[:,2], color='green', label=r'$Wasserstein$', linestyle='dashed')
+        plt.plot(epoch, loss[:,3], color='royalblue', markersize=12, label=r'$Gradient$', linestyle='dashed')
+        plt.plot(epoch, Rec_scaling*loss[:,4], color='orange', markersize=12, label=r'$Reconstrucion$', linestyle='dashed')
+    except:
+        plt.plot(epoch, loss[:,0], color='red', markersize=12, label=r'$Total$')
+
+    plt.legend(loc='upper right',prop={'size':10})
+    ax1.set_xlabel(r'$Iterations$')
+    ax1.set_ylabel(r'$Cost \Theta (MSE)$')
+    ax1.set_ylim(-2, 2)
+    fig.savefig(log_dir + '/%s_Loss.png' %name, dpi=480)
+    plt.close('all')
+
+
+def plot_array_scatter(values, label, axis, vmin=0):
+    """Plot a map *values* for an detector array specified by *v_stations*. """
+    xd, yd = rectangular_array()
+    filter = values!=0
+    filter[5,5] = True
+    axis.scatter(xd[~filter], yd[~filter], c = 'grey', s = 220, alpha=0.1, label="silent")
+    circles = axis.scatter(xd[filter], yd[filter], c = values[filter], s = 200, alpha=1, label="loud", vmin=vmin)
+    cbar = plt.colorbar(circles, ax=axis)
+    cbar.set_label(label + ' [a.u.]')
+    axis.grid(True)
+    axis.set_title(label + ' footprint', fontsize=10)
+    axis.set_aspect('equal')
+    axis.set_xlim(-5, 5)
+    axis.set_ylim(-5, 5)
+    axis.set_xlabel('x [km]')
+    axis.set_ylabel('y [km]')
+
+
+def plot_showermaps(footprint, log_dir='.', title='', epoch=''):
+    """ Plots the time and signal footprint in one figure """
+    fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, figsize=(10,4))
+    fig.subplots_adjust(left=0, bottom=0.13)
+    plot_array_scatter(footprint[:,:,0], label = 'time', axis=ax1, vmin=1)
+    plot_array_scatter(footprint[:,:,1], label = 'signal', axis=ax2)
+    plt.suptitle(title + ' ' + str(epoch), fontsize=16)
+    plt.tight_layout()
+    plt.legend(bbox_to_anchor=(0.59, 0.12, 0, 0.0), bbox_transform=plt.gcf().transFigure, ncol=2)
+    fig.savefig(log_dir + '/%s_shower_map.png' %epoch, dpi=140)
+    plt.close('all')
+
+
+def plot_signal_map(footprint, axis, label=None):
+    """Plot a map *footprint* for an detector array specified by *v_stations*. """
+    xd, yd = rectangular_array()
+    filter = footprint!=0
+    filter[5,5] = True
+    axis.scatter(xd[~filter], yd[~filter], c = 'grey', s = 150, alpha=0.1, label="silent")
+    circles = axis.scatter(xd[filter], yd[filter], c = footprint[filter], s = 150, alpha=1, label="loud")
+    cbar = plt.colorbar(circles, ax=axis)
+    cbar.set_label('signal [a.u.]')
+    axis.grid(True)
+    if label!=None:
+        axis.text(0.95, 0.1, "Energy: %.1f EeV" %label, verticalalignment = 'top', horizontalalignment = 'right', transform=axis.transAxes, backgroundcolor='w')
+    axis.set_aspect('equal')
+    axis.set_xlim(-5, 5)
+    axis.set_ylim(-5, 5)
+    axis.set_xlabel('x [km]')
+    axis.set_ylabel('y [km]')
+
+
+def plot_multiple_signalmaps(footprint, log_dir='.', title='', epoch='', nrows=2, ncols=2, labels=None, shuffle=True):
+    """ Plots the time and signal footprint in one figure """
+    if shuffle:
+        idx = np.random.choice(np.arange(footprint.shape[0]), size=nrows*ncols)
+    else:
+        idx = np.arange(footprint.shape[0])
+    fig, sub = plt.subplots(nrows=nrows, ncols=ncols, figsize=(9,7))
+    for i in range(ncols):
+        for j in range(nrows):
+            try:
+                plot_signal_map(footprint[idx[ncols*i+j]], axis=sub[i,j], label=labels[idx[ncols*i+j]])
+            except:
+                plot_signal_map(footprint[idx[ncols*i+j]], axis=sub[i,j], label=labels)
+    plt.tight_layout()
+    fig.subplots_adjust(left=0.02, top=0.95)
+    plt.suptitle(title + ' ' + str(epoch), fontsize=12)
+    fig.savefig(log_dir + '/%s_signal_maps.png' %epoch, dpi=140)
+    plt.close('all')
+
+
+def plot_energy_histos(y_true, y_pred, log_dir, name):
+    ''' Plottingfunction to evaluate the reconstruction of the cosmic ray energy
+        plots: predicted Energy against MC Energy and the energy dependence of the relative energy resolution (embedded)'''
+    y_true = y_true.reshape(y_true.shape[0],1)
+    reco = y_true-y_pred
+    # Embedding plot
+    fig, ax1 = plt.subplots(1)
+    ax1.scatter(y_true, y_pred, s = 2, color='royalblue', alpha=.90)
+    ax1.set_xlabel(r'$Energy_{Input}\;[EeV]$')
+    ax1.set_ylabel(r'$Energy_{DNN}\;[EeV]$')
+    left, bottom, width, height = [0.23, 0.6, 0.3, 0.26] # Links Oben
+    ax1.text(0.95,0.2, r"$\mu = $ %.3f" % np.mean(reco) + " [EeV]" + "\n" + "$\sigma = $ %.3f" % np.std(reco) + " [EeV]", verticalalignment = 'top', horizontalalignment = 'right', transform=ax1.transAxes, backgroundcolor='w')
+    ax1.plot([np.min(y_true), np.max(y_true)], [np.min(y_true), np.max(y_true)], color='red')
+    fig.savefig(log_dir + '/%s_Energy_reconstruction.png' %name, dpi=140)
+    return reco
+
+def plot_energy_histos_real(y_true, y_pred, log_dir, name):
+    ''' Plottingfunction to evaluate the reconstruction of the cosmic ray energy
+        plots: predicted Energy against MC Energy and the energy dependence of the relative energy resolution (embedded)'''
+    y_true = y_true.reshape(y_true.shape[0],1)
+    reco = y_true-y_pred
+    # Embedding plot
+    fig, ax1 = plt.subplots(1)
+    ax1.scatter(y_true, y_pred, s = 2, color='royalblue', alpha=.90)
+    ax1.set_xlabel(r'$Energy_{MC}\;[EeV]$')
+    ax1.set_ylabel(r'$Energy_{DNN}\;[EeV]$')
+    left, bottom, width, height = [0.23, 0.6, 0.3, 0.26] # Links Oben
+    ax1.text(0.95,0.2, r"$\mu = $ %.3f" % np.mean(reco) + " [EeV]" + "\n" + "$\sigma = $ %.3f" % np.std(reco) + " [EeV]", verticalalignment = 'top', horizontalalignment = 'right', transform=ax1.transAxes, backgroundcolor='w')
+    ax1.plot([np.min(y_true), np.max(y_true)], [np.min(y_true), np.max(y_true)], color='red')
+    fig.savefig(log_dir + '/%s_Energy_reconstruction.png' %name, dpi=140)
+    return reco
+