}
colnames = [
r"$\sum_{i} E_{Ti}$ [GeV]", r"$\max (E_{Ti})$ [GeV]",
r"$\sum_{i} [E_{Ti} > 6MeV]$", r"std($E_{Ti}$) [GeV]",
r"$E_{T,Tracker} - \sum_{i} E_{Ti}$ [MeV]"
]
for func_idx, (func, params) in enumerate(use_functions.items()):
if func.__name__ in [
"get_number_of_activated_cells", "get_max_energy"
]:
build_histogram(true=mc_data_images_m,
fake=generated_images,
function=func,
name=colnames[func_idx],
epoch="",
folder=None,
ax=axes[func_idx],
labels=["Geant4", "Im2Im"],
**params)
else:
build_histogram(true=mc_data_images_m,
fake=generated_images,
function=func,
name=colnames[func_idx],
epoch="",
folder=None,
ax=axes[func_idx],
labels=["Geant4", "Im2Im"],
**params)
generated_images = generated_images.reshape([-1, *image_shape])
fig2, axes = plt.subplots(3, 3, figsize=(20, 20))
axes = np.ravel(axes)
use_functions = {get_energies: {"energy_scaler": 1}, get_max_energy: {"energy_scaler": 1, "maxclip": 6.120},
get_number_of_activated_cells: {"threshold": 5/1000},
get_center_of_mass_x: {"image_shape": image_shape}, get_center_of_mass_y: {"image_shape": image_shape},
get_std_energy: {"energy_scaler": 1}, get_energy_resolution: {"real_ET": tracker_real_ET/1000, "energy_scaler": 1}}
colnames = ["Enery", "MaxEnergy", "Cells", "X CoM", "Y CoM", "StdEnergy", "Resolution", "dN/dE", "(dN/dE) / (dN/dE)"]
htos_calo_images_real = crop_images(calo_images, **padding)
htos_calo_images_fake = crop_images(generated_images, **padding)
for func_idx, (func, params) in enumerate(use_functions.items()):
if func.__name__ in ["get_number_of_activated_cells", "get_max_energy"]:
build_histogram(true=htos_calo_images_real, fake=htos_calo_images_fake, function=func,
name=colnames[func_idx], epoch="", folder=None, ax=axes[func_idx], **params)
else:
build_histogram(true=calo_images, fake=generated_images, function=func,
name=colnames[func_idx], epoch="", folder=None, ax=axes[func_idx], **params)
build_histogram_HTOS(true=htos_calo_images_real, fake=htos_calo_images_fake,
energy_scaler=1, threshold=3.6, real_ET=tracker_real_ET,
ax1=axes[7], ax2=axes[8])
figs.append(fig2)
for idx in range(10):
print(idx+1, "/", 10)
use_functions = {
get_energies: {"energy_scaler": 1}, get_max_energy: {"energy_scaler": 1, "maxclip": 6.120},
get_number_of_activated_cells: {"threshold": 5/1000},
get_center_of_mass_x: {"image_shape": image_shape}, get_center_of_mass_y: {"image_shape": image_shape},
def train(d_model, g_model, gan_model, dataset, n_epochs=100, n_batch=1):
trainA, trainB = dataset
batches_per_epo = int(len(trainA) / n_batch)
n_steps = batches_per_epo * n_epochs
for i in range(n_steps):
n_patch = d_model.layers[-1].output_shape[1]
[X_realA,
X_realB], y_real = generate_real_samples(dataset, n_batch, n_patch)
X_fakeB, y_fake = generate_fake_samples(g_model, X_realA, n_patch)
d_loss1 = d_model.train_on_batch([X_realA, X_realB], y_real)
d_loss2 = d_model.train_on_batch([X_realA, X_fakeB], y_fake)
g_loss, _, _ = gan_model.train_on_batch(X_realA, [y_real, X_realB])
if i % 5000 == 0:
print('>%d / %d, d1[%.3f] d2[%.3f] g[%.3f]' %
(i + 1, n_steps, d_loss1, d_loss2, g_loss))
nr_eval = 300
im_shape = [64, 64]
samples_y = g_model.predict(test_x[:nr_eval]).reshape(
[-1, *im_shape])
samples_x = test_x[:nr_eval].reshape([-1, *im_shape])
array_of_images = np.array([[
tx.reshape(image_shape[0:2]),
ty.reshape(image_shape[0:2]),
gx.reshape(image_shape[0:2])
] for tx, ty, gx in zip(test_x[:10], test_y[:10], samples_y[:10])])
fig, ax = build_images(array_of_images,
column_titles=["X", "Y", "G(X)"],
fs_x=45,
fs_y=25)
plt.savefig("{}/GeneratedSamples/Samples_{}.png".format(
path_saving, i))
build_histogram(true=samples_x,
fake=samples_y,
function=get_energies,
name="Energy",
epoch=i,
folder=path_saving,
energy_scaler=energy_scaler)
build_histogram(true=samples_x,
fake=samples_y,
function=get_number_of_activated_cells,
name="Cells",
epoch=i,
folder=path_saving,
threshold=5 / energy_scaler)
build_histogram(true=samples_x,
fake=samples_y,
function=get_max_energy,
name="MaxEnergy",
epoch=i,
folder=path_saving,
energy_scaler=energy_scaler)
build_histogram(true=samples_x,
fake=samples_y,
function=get_center_of_mass_x,
name="CenterOfMassX",
epoch=i,
folder=path_saving,
image_shape=im_shape)
build_histogram(true=samples_x,
fake=samples_y,
function=get_center_of_mass_y,
name="CenterOfMassY",
epoch=i,
folder=path_saving,
image_shape=im_shape)
build_histogram(true=samples_x,
fake=samples_y,
function=get_std_energy,
name="StdEnergy",
epoch=i,
folder=path_saving,
energy_scaler=energy_scaler)
save_files = os.listdir(path_saving + "/ModelSave")
for f in save_files:
os.remove(path_saving + "/ModelSave/" + f)
with open(path_saving + "/ModelSave/Generator_{}.pickle".format(i),
"wb") as f:
pickle.dump(g_model, f)
with open(
path_saving +
"/ModelSave/Discriminator_{}.pickle".format(i), "wb") as f:
pickle.dump(d_model, f)
with open(path_saving + "/ModelSave/Model_{}.pickle".format(i),
"wb") as f:
pickle.dump(gan_model, f)
def evaluate(self, true, condition, epoch):
print("Batch ", epoch)
log_start = time.clock()
self._batches.append(epoch)
fake = self._sess.run(self._output_gen_from_encoding,
feed_dict={
self._X_input: self._x_test,
self._Z_input: self._z_test,
self._is_training: False
})
true = self._y_test.reshape(
[-1, self._image_shape[0], self._image_shape[1]])
fake = fake.reshape([-1, self._image_shape[0], self._image_shape[1]])
build_histogram(true=true,
fake=fake,
function=get_energies,
name="Energy",
epoch=epoch,
folder=self._folder)
build_histogram(true=true,
fake=fake,
function=get_number_of_activated_cells,
name="Cells",
epoch=epoch,
folder=self._folder,
threshold=6 / 6120)
build_histogram(true=true,
fake=fake,
function=get_max_energy,
name="MaxEnergy",
epoch=epoch,
folder=self._folder)
build_histogram(true=true,
fake=fake,
function=get_center_of_mass_x,
name="CenterOfMassX",
epoch=epoch,
folder=self._folder,
image_shape=self._image_shape)
build_histogram(true=true,
fake=fake,
function=get_center_of_mass_y,
name="CenterOfMassY",
epoch=epoch,
folder=self._folder,
image_shape=self._image_shape)
build_histogram(true=true,
fake=fake,
function=get_std_energy,
name="StdEnergy",
epoch=epoch,
folder=self._folder)
fig, axs = plt.subplots(nrows=2, ncols=2, figsize=(18, 10))
axs = np.ravel(axs)
if "Gradients" in self._monitor_dict:
colors = ["green", "blue", "red"]
axy_min = np.inf
axy_max = -np.inf
for go, gradient_ops in enumerate(
self._monitor_dict["Gradients"][0]):
grads = [
self._sess.run(
gv[0],
feed_dict={
self._X_input:
self._x_test,
self._Y_input:
self._y_test,
self._Z_input:
self._z_test,
self._is_training:
False,
self._output_label_real:
self.get_random_label(is_real=True,
size=self._nr_test),
self._output_label_fake:
self.get_random_label(is_real=False,
size=self._nr_test)
}) for gv in gradient_ops
]
for op_idx, op in enumerate([np.mean, np.max, np.min]):
self._monitor_dict["Gradients"][2][go * 3 + op_idx].append(
op([op(grad) for grad in grads]))
vals = self._monitor_dict["Gradients"][2][go * 3 + op_idx]
if op_idx == 0:
axs[0].plot(
self._batches,
vals,
label=self._monitor_dict["Gradients"][1][go],
color=colors[go])
else:
axs[0].plot(self._batches,
vals,
linewidth=0.5,
linestyle="--",
color=colors[go])
upper = np.mean(vals)
lower = np.mean(vals)
if upper > axy_max:
axy_max = upper
if lower < axy_min:
axy_min = lower
axs[0].set_title("Gradients")
axs[0].legend()
axs[0].set_ylim([axy_min, axy_max])
current_batch_x, current_batch_y = self._trainset.get_next_batch(
self.batch_size)
Z_noise = self._generator.sample_noise(n=len(current_batch_x))
colors = [
"green", "blue", "red", "orange", "purple", "brown", "gray",
"pink", "cyan", "olive"
]
for k, key in enumerate(self._monitor_dict):
if key == "Gradients":
continue
key_results = self._sess.run(
self._monitor_dict[key][0],
feed_dict={
self._X_input: current_batch_x,
self._Y_input: current_batch_y,
self._Z_input: Z_noise,
self._is_training: True,
self._output_label_real:
self.get_random_label(is_real=True),
self._output_label_fake:
self.get_random_label(is_real=False)
})
for kr, key_result in enumerate(key_results):
try:
self._monitor_dict[key][2][kr].append(
self._monitor_dict[key][3][0](key_result))
except IndexError:
self._monitor_dict[key][2][kr].append(key_result)
axs[k].plot(self._batches,
self._monitor_dict[key][2][kr],
label=self._monitor_dict[key][1][kr],
color=colors[kr])
axs[k].legend()
axs[k].set_title(key)
print("; ".join([
"{}: {}".format(name, round(float(val[-1]), 5)) for name, val
in zip(self._monitor_dict[key][1], self._monitor_dict[key][2])
]))
gen_samples = self._sess.run(
[self._output_gen_from_encoding],
feed_dict={
self._X_input: current_batch_x,
self._Z_input: Z_noise,
self._is_training: False
})
axs[-1].hist([np.ravel(gen_samples),
np.ravel(current_batch_y)],
label=["Generated", "True"])
axs[-1].set_title("Pixel distribution")
axs[-1].legend()
for check_idx, check_key in enumerate(self._check_dict):
result_bools_of_check = []
check = self._check_dict[check_key]
for tensor_idx in range(len(check["Tensors"])):
tensor_ = self._sess.run(check["Tensors"][tensor_idx],
feed_dict={
self._X_input: self._x_test,
self._Y_input: self._y_test,
self._Z_input: self._z_test,
self._is_training: False
})
tensor_op = check["OPonTensors"][tensor_idx](tensor_)
if eval(
str(tensor_op) + check["Relation"][tensor_idx] +
str(check["Threshold"][tensor_idx])):
result_bools_of_check.append(True)
else:
result_bools_of_check.append(False)
if (tensor_idx > 0 and check["TensorRelation"](
*result_bools_of_check)) or (result_bools_of_check[0]):
self._check_count[check_idx] += 1
if self._check_count[
check_idx] == self._max_allowed_failed_checks:
raise GeneratorExit(check_key)
else:
self._check_count[check_idx] = 0
self._total_log_time += (time.clock() - log_start)
fig.suptitle("Train {} / Log {} / Fails {}".format(
np.round(self._total_train_time, 2),
np.round(self._total_log_time, 2), self._check_count))
plt.savefig(self._folder + "/TrainStatistics.png")
plt.close("all")
r'$\sigma=%.2f$' % (resolution_std, )))
props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
axes[model_idx, -4].text(0.05,
0.95,
textstr,
transform=axes[model_idx, -4].transAxes,
fontsize=14,
verticalalignment='top',
bbox=props)
for func_idx, (func, params) in enumerate(use_functions.items()):
build_histogram(true=mc_data_images_m,
fake=generated_images_from_cgan,
fake2=generated_images_from_tracker,
function=func,
name=colnames[func_idx],
epoch="",
folder=None,
ax=axes[model_idx, func_idx],
labels=["Geant4", "CGAN", "Tracker"],
**params)
if func_idx == 0:
fs = {
"2": 12,
"3": 15,
"7": 10,
"8": 10,
"9": 12,
"11": 13,
"12": 12,
"13": 12,
"14": 13,
}
build_histogram_HTOS(true=calo_images,
fake=None,
energy_scaler=calo_scaler,
threshold=3600,
real_ET=tracker_real_ET,
ax1=axes[7],
ax2=axes[8])
for func_idx, (func, params) in enumerate(use_functions.items()):
build_histogram(true=calo_images,
fake=None,
function=func,
name=colnames[func_idx],
epoch="",
folder=None,
ax=axes[func_idx],
use_legend=True,
**params)
figs.append(fig_caloStat)
fig_calo_distribution, ax = plt.subplots(nrows=1,
ncols=3,
figsize=(12, 8),
facecolor='w',
edgecolor='k')
fig_calo_distribution.subplots_adjust(wspace=0.2, hspace=0.05)
## Distribution total energy
def train(self,
x_train,
y_train,
x_test=None,
y_test=None,
epochs=100,
batch_size=32,
evaluate_interval=50,
label_smoothing=1,
adversarial_steps=1):
out_patch = self.critic.layers[-1].output_shape[1:]
real_labels = -np.ones((batch_size, *out_patch)) * label_smoothing
fake_labels = np.ones((batch_size, *out_patch)) * label_smoothing
dummy_labels = np.zeros(
(batch_size,
*out_patch)) * label_smoothing # Dummy gt for gradient penalty
self._gen_out_zero = 0
for epoch in range(epochs):
for i in range(adversarial_steps):
# ---------------------
# Train Discriminator
# ---------------------
# Select a random batch of images
idx = np.random.randint(0, x_train.shape[0], batch_size)
train_images = x_train[idx]
train_labels = y_train[idx]
train_labels = np.array([
np.tile(lbl, (*image_shape[:-1], 1))
for lbl in train_labels
])
# Sample generator input
train_noise = np.random.normal(0, 1,
(batch_size, self.latent_dim))
train_input = np.concatenate([train_noise, y_train[idx]],
axis=1)
# Train the critic
d_loss = self.critic_model.train_on_batch(
[train_images, train_input, train_labels],
[real_labels, fake_labels, dummy_labels])
# ---------------------
# Train Generator
# ---------------------
g_loss = self.generator_model.train_on_batch(
[train_input, train_labels], real_labels)
if np.isinf(g_loss) or np.isnan(g_loss) or np.isnan(
d_loss[0]) or np.isinf(d_loss[0]):
raise GeneratorExit(
"Nan or inf in losses: d_loss: {} |-| g_loss: {}.".format(
d_loss, g_loss))
# Print the progress
if epoch % evaluate_interval == 0:
print(
">%d / %d, d_tot[%.3f], d_real[%.3f], d_fake[%.3f], d_grad[%.3f], g[%.3f]"
% (epoch + 1, epochs, d_loss[0], d_loss[1], d_loss[2],
d_loss[3], g_loss))
nr_eval = 100
im_shape = self.image_shape
test_noise = np.random.normal(0, 1, (nr_eval, self.latent_dim))
test_input = np.concatenate([test_noise, y_test[:nr_eval]],
axis=1)
samples_gen = self.generator.predict(test_input).reshape(
[-1, *im_shape[:-1]])
samples_x = x_test[:nr_eval].reshape([-1, *im_shape[:-1]])
array_of_images = np.array(
[[tx.reshape(im_shape[0:2]),
gx.reshape(im_shape[0:2])]
for tx, gx in zip(x_test[:10], samples_gen[:10])])
fig, ax = build_images(array_of_images,
column_titles=["X", "Y", "G(X)"],
fs_x=45,
fs_y=25)
plt.savefig("{}/GeneratedSamples/Samples_{}.png".format(
path_saving, epoch))
energy_scaler = scaler["Calo"]
build_histogram(true=samples_x,
fake=samples_gen,
function=get_energies,
name="Energy",
epoch=epoch,
folder=path_saving,
energy_scaler=energy_scaler)
build_histogram(true=samples_x,
fake=samples_gen,
function=get_number_of_activated_cells,
name="Cells",
epoch=epoch,
folder=path_saving,
threshold=5 / energy_scaler)
build_histogram(true=samples_x,
fake=samples_gen,
function=get_max_energy,
name="MaxEnergy",
epoch=epoch,
folder=path_saving,
energy_scaler=energy_scaler)
build_histogram(true=samples_x,
fake=samples_gen,
function=get_center_of_mass_x,
name="CenterOfMassX",
epoch=epoch,
folder=path_saving,
image_shape=im_shape)
build_histogram(true=samples_x,
fake=samples_gen,
function=get_center_of_mass_y,
name="CenterOfMassY",
epoch=epoch,
folder=path_saving,
image_shape=im_shape)
build_histogram(true=samples_x,
fake=samples_gen,
function=get_std_energy,
name="StdEnergy",
epoch=epoch,
folder=path_saving,
energy_scaler=energy_scaler)
save_files = os.listdir(path_saving + "/ModelSave")
for f in save_files:
os.remove(path_saving + "/ModelSave/" + f)
with open(
path_saving +
"/ModelSave/Generator_{}.pickle".format(epoch),
"wb") as f:
pickle.dump(self.generator, f)
with open(
path_saving +
"/ModelSave/Discriminator_{}.pickle".format(epoch),
"wb") as f:
pickle.dump(self.critic, f)
self.sample_images(epoch, samples_gen)
samples_gen = 0
if np.max(samples_gen) < 0.05:
self._gen_out_zero += 1
if self._gen_out_zero == 20:
raise GeneratorExit("Generator outputs zeros")
else:
self._gen_out_zero = 0