diff --git a/architectures/bicycle_GAN.py b/architectures/bicycle_GAN.py index 0b2bf3c..f580ccb 100644 --- a/architectures/bicycle_GAN.py +++ b/architectures/bicycle_GAN.py @@ -8,10 +8,28 @@ from architectures.utils.NN_building_blocks import * from architectures.utils.NN_gen_building_blocks import * +from architectures.utils.arbitrary_act import * def lrelu(x, alpha=0.2): return tf.maximum(alpha*x,x) +# def tf_HCAL_act(x, mask): + +# output = tf.where(mask, +# tf.multiply(6.120,tf.nn.sigmoid(tf.subtract(x,5.))), +# tf.multiply(1.530,tf.nn.sigmoid(tf.subtract(x,5.))) +# ) + +# return output + +def tf_HCAL_act(x, mask): + + output = tf.where(mask, + tf.maximum(0.,tf.minimum(6120.,x)) , + tf.maximum(0.,tf.minimum(1530.,x)) + ) + + return output pretrain=None # some dummy constants @@ -169,6 +187,12 @@ name='X_test_A', ) + self.mask = tf.placeholder( + tf.bool, + shape=(None, + n_H_A, n_W_A, n_C), + name='inner_outer', + ) G = bicycleGenerator(self.input_A, self.n_H_B, self.n_W_B, g_sizes_enc, g_sizes_dec, 'A_to_B') @@ -181,11 +205,11 @@ z_encoded, z_encoded_mu, z_encoded_log_sigma = E.e_forward(self.input_B) with tf.variable_scope('generator_A_to_B') as scope: - sample_A_to_B_encoded = G.g_forward(self.input_A, z_encoded) + sample_A_to_B_encoded = tf_HCAL_act(G.g_forward(self.input_A, z_encoded), self.mask) with tf.variable_scope('generator_A_to_B') as scope: scope.reuse_variables() - sample_A_to_B = self.sample_A_to_B = G.g_forward(self.input_A, self.z, reuse=True) + sample_A_to_B = self.sample_A_to_B = tf_HCAL_act(G.g_forward(self.input_A, self.z, reuse=True), self.mask) with tf.variable_scope('encoder_B') as scope: scope.reuse_variables() @@ -203,9 +227,9 @@ with tf.variable_scope('generator_A_to_B') as scope: scope.reuse_variables() - self.test_images_A_to_B = G.g_forward( + self.test_images_A_to_B = tf_HCAL_act(G.g_forward( self.input_test_A, self.z, reuse=True, is_training=False - ) + ),self.mask) #parameters lists self.d_params =[t for t in tf.trainable_variables() if 'discriminator' in t.name] @@ -221,7 +245,7 @@ predicted_fake_encoded = tf.nn.sigmoid(logits_fake_encoded) #predicted_fake_encoded =tf.maximum(tf.minimum(predicted_fake_encoded, 0.99), 0.00) - epsilon=1e-3 + epsilon=1e-2 #GAN LOSS if cost_type=='GAN': @@ -298,6 +322,7 @@ labels=epsilon+tf.zeros_like(logits_fake_encoded) ) ) + #GENERATOR LOSSES g_cost_lr_GAN = tf.reduce_sum(tf.squared_difference(feature_output_real,feature_output_fake), axis=1) @@ -308,7 +333,8 @@ #CYCLIC WEIGHT - self.g_cost_cycl = tf.reduce_mean(tf.abs(self.input_B - sample_A_to_B_encoded)) + g_cost_cycl = tf.reduce_mean(tf.abs(self.input_B - sample_A_to_B_encoded), axis=[1,2,3]) + self.g_cost_cycl = tf.reduce_mean(g_cost_cycl) # self.g_4_cells_cycl_encoded=tf.reduce_mean( @@ -352,14 +378,16 @@ #ENCODER COSTS - self.e_cost_latent_cycle = tf.reduce_mean(tf.abs(self.z - z_recon)) + e_cost_latent_cycle = tf.reduce_mean(tf.abs(self.z - z_recon), axis=1) + self.e_cost_latent_cycle=tf.reduce_mean(e_cost_latent_cycle) + self.e_cost_kl = -0.5 * tf.reduce_mean(1 + 2*tf.log(z_encoded_log_sigma) - tf.square(z_encoded_mu) - tf.square(z_encoded_log_sigma)) #TOTAL COSTS self.d_cost = self.d_cost_fake_vae_GAN + self.d_cost_fake_lr_GAN + 2*self.d_cost_real - self.g_cost = self.g_cost_vae_GAN + self.g_cost_lr_GAN + cycl_weight*(self.g_cost_cycl) + self.e_cost_latent_cycle - self.e_cost = latent_weight*self.e_cost_latent_cycle + kl_weight*self.e_cost_kl + self.g_cost_vae_GAN + self.g_cost_cycl + self.g_cost = self.g_cost_vae_GAN + self.g_cost_lr_GAN + cycl_weight*self.g_cost_cycl + self.e_cost_latent_cycle + self.e_cost = latent_weight*self.e_cost_latent_cycle + kl_weight*self.e_cost_kl + self.g_cost_vae_GAN + self.g_cost_cycl self.d_train_op = tf.train.AdamOptimizer( learning_rate=lr, @@ -491,6 +519,10 @@ bs=X_batch_A.shape[0] + in_out_mask_zeros = np.zeros([bs, self.n_H_A, self.n_W_A, self.n_C]) + in_out_mask_zeros[:, 12:40, 16:48, :]=1 + in_out_mask = in_out_mask_zeros.astype(dtype=bool) + t0 = datetime.now() @@ -508,17 +540,17 @@ d_cost_vae_GAN=0 d_cost_lr_GAN=0 - + sample_z = np.random.normal(size=(bs, self.latent_dims)) for i in range(discr_steps): - sample_z = np.random.normal(size=(bs, self.latent_dims)) + #sample_z = np.random.normal(size=(bs, self.latent_dims)) _, d_acc, d_acc_enc, d_cost, d_cost_vae_GAN, d_cost_lr_GAN = self.session.run( (self.d_train_op, self.d_accuracy, self.d_accuracy_enc, self.d_cost, self.d_cost_fake_vae_GAN, self.d_cost_fake_lr_GAN), feed_dict={self.input_A:X_batch_A, self.input_B:X_batch_B, - self.z:sample_z, self.batch_sz:bs + self.z:sample_z, self.mask:in_out_mask, self.batch_sz:bs }, ) d_cost+=d_cost @@ -531,7 +563,7 @@ for i in range(gen_steps): - sample_z = np.random.normal(size=(bs, self.latent_dims)) + #sample_z = np.random.normal(size=(bs, self.latent_dims)) _, g_cost, g_cost_cycl, g_cost_lr_GAN, g_cost_vae_GAN = self.session.run( @@ -540,7 +572,7 @@ ), feed_dict={self.input_A:X_batch_A, self.input_B:X_batch_B, - self.z:sample_z, self.batch_sz:bs + self.z:sample_z, self.mask:in_out_mask, self.batch_sz:bs }, ) @@ -558,8 +590,8 @@ #cluster_diffs.append(self.cycl_weight*cluster_diff/gen_steps) for i in range(vae_steps): - sample_z = np.random.normal(size=(bs, self.latent_dims)) - + + #sample_z = np.random.normal(size=(bs, self.latent_dims)) _, e_cost, e_cost_latent_cycle, e_cost_kl = self.session.run( (self.e_train_op, self.e_cost, @@ -567,7 +599,7 @@ ), feed_dict={self.input_A:X_batch_A, self.input_B:X_batch_B, - self.z:sample_z, self.batch_sz:bs + self.z:sample_z, self.mask:in_out_mask, self.batch_sz:bs }, ) @@ -677,6 +709,8 @@ validation_true_hist = np.sum(validation_true_hist,axis=1) max_true_hist = np.max(validation_true_hist) + validation_NN_hist_rescaled=(validation_NN_hist/max_NN_hist)*max_MC_hist + else: validation_MC_hist= validating_B.reshape(validating_size, self.n_H_B*self.n_W_B) validation_MC_hist = np.sum(validation_MC_hist,axis=1) @@ -689,7 +723,10 @@ validation_true_hist= validating_A.reshape(validating_size, self.n_H_A*self.n_W_A) validation_true_hist = np.sum(validation_true_hist,axis=1) max_true_hist = np.max(validation_true_hist) - + + validation_NN_hist_rescaled=(validation_NN_hist/max_NN_hist)*max_MC_hist + + plt.clf() plt.subplot(1,3,1) h_reco = plt.hist(validation_true_hist,bins=30, edgecolor='black'); @@ -702,7 +739,7 @@ plt.ylabel('dN/dE') plt.title('Reco pion E_T distribution,\n max MC hist: {0} '.format(max_MC_hist)) plt.subplot(1,3,3) - h_nn = plt.hist(validation_NN_hist,bins=30, edgecolor='black'); + h_nn = plt.hist(validation_NN_hist_rescaled,bins=30, edgecolor='black'); plt.xlabel('E (MeV)') plt.ylabel('dN/dE') plt.title('Reco pion E_T distribution from bicycleGAN, \n max NN hist: {0} '.format(max_NN_hist)) @@ -725,17 +762,17 @@ print('Resolution plots are being printed...') - diffNN = (validation_NN_hist/max_NN_hist)*max_MC_hist-validation_true_hist + diffNN = validation_NN_hist_rescaled-validation_true_hist diffMC = validation_MC_hist-validation_true_hist plt.clf() plt.subplot(1,2,1) - h_reco = plt.hist(diffMC,bins=30, edgecolor='black'); + h_reco = plt.hist(diffMC/1000,bins=30, range=(-80, 40), edgecolor='black'); plt.xlabel('ET recoMC - ET true') plt.ylabel('dN/dETdiff') plt.title('Resolution as simulated by MC') plt.subplot(1,2,2) - h_nn = plt.hist(diffNN,bins=30, edgecolor='black'); + h_nn = plt.hist(diffNN/1000,bins=30, range=(-80, 40), edgecolor='black'); plt.xlabel('ET recoNN - ET true') plt.ylabel('dN/dETdiff') plt.title('Resolution as simulated by NN') @@ -746,23 +783,30 @@ def get_sample_A_to_B(self, X): - z = np.random.normal(size=(1, self.latent_dims)) + in_out_mask_zeros = np.zeros([1, self.n_H_A, self.n_W_A, self.n_C]) + in_out_mask_zeros[:, 12:40, 16:48, :]=1 + in_out_mask = in_out_mask_zeros.astype(dtype=bool) + one_sample = self.session.run( self.test_images_A_to_B, - feed_dict={self.input_test_A:X, self.z:z, self.batch_sz: 1}) + feed_dict={self.input_test_A:X, self.z:z, self.mask:in_out_mask, self.batch_sz: 1}) return one_sample def get_samples_A_to_B(self, X): - - z = np.random.normal(size=(X.shape[0], self.latent_dims)) - + bs=X.shape[0] + z = np.random.normal(size=(bs, self.latent_dims)) + + in_out_mask_zeros = np.zeros([bs, self.n_H_A, self.n_W_A, self.n_C]) + in_out_mask_zeros[:, 12:40, 16:48, :]=1 + in_out_mask = in_out_mask_zeros.astype(dtype=bool) + many_samples = self.session.run( self.test_images_A_to_B, - feed_dict={self.input_test_A:X, self.z:z, self.batch_sz: X.shape[0]}) + feed_dict={self.input_test_A:X, self.z:z, self.mask:in_out_mask, self.batch_sz: X.shape[0]}) return many_samples