Skip to content
Snippets Groups Projects
Commit af5bae04 authored by Krishna Krishna Nikhil's avatar Krishna Krishna Nikhil
Browse files

Edited Reco Model

parent 839ed40a
Branches
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
Showing
with 192 additions and 58 deletions
modules/execute.py
+ 192
− 58
View file @ af5bae04
import multiprocessing
import pickle
from generator import Generator, Parameters, SurrogateDataset
from reconstruction import Reconstruction
......@@ -9,12 +8,27 @@ from optimizerA import Optimizer
from matplotlib import pyplot as plt
import time
import torch
from torch.utils.data import DataLoader
import numpy as np
import os
import json
import sys
import wandb
import argparse
import random
# Parse command-line arguments
parser = argparse.ArgumentParser()
parser.add_argument('--seed', type=int, required=True, help='Random seed')
args = parser.parse_args()
print(args.seed)
# Set the random seed
random.seed(args.seed)
os.nice(18)
# Login to wandb using API key
wandb.login(key='1b3bc8141d462316caf2c4bb693ae40aa34f8bd6')
def reset_weights(m):
'''
......@@ -28,53 +42,65 @@ def reset_weights(m):
if __name__ == "__main__":
multiprocessing.set_start_method('spawn')
outpath = 'FF04'
run_num = (args.seed / 100) + 1 # + 1 for just this run brother
outpath = 'Fixing' + str(run_num)
os.system('mkdir '+outpath)
outpath += '/'
os.system('cp *.py '+outpath)
os.system('cp FinalA.py '+outpath)
os.system('cp cleaning.py '+outpath)
os.system('cp plotting.py '+outpath)
os.system('cp optimizerA.py '+outpath)
test = False
divide = 10 if test else 1
num_layers = 3
#so far these are only layer thicknesses
start_pars = {'thickness_absorber_0': 35,
'thickness_absorber_1': 17,
'thickness_absorber_2': 20,
#'thickness_absorber_5': .1,
#'thickness_absorber_6': .1,
#'thickness_absorber_7': .1,
#'thickness_absorber_8': .1,
'thickness_scintillator_0': 5,
'thickness_scintillator_1': 12,
'thickness_scintillator_2': 3,
#'thickness_scintillator_5': 0.5,
#'thickness_scintillator_6': 0.5,
#'thickness_scintillator_7': 0.5,
#'thickness_scintillator_8': 0.5,
'material_abs_0' : 0.33,
'material_abs_1' : 0.33,
'material_abs_2' : 0.33, # Randomly chosen values see what happens, ideally we could pass an array here, this for testing
'material_scint_0' : 0.33,
'material_scint_1' : 0.66,
'material_scint_2' : 0.33,
}
ToP = 2
import random
start_pars = {
'thickness_absorber_0': random.uniform(0.1, 50), # Random float between 0.1 and 50
'thickness_absorber_1': random.uniform(0.1, 50),
'thickness_absorber_2': random.uniform(0.1, 50),
'thickness_scintillator_0': random.uniform(0.1, 50), # Random float between 0.1 and 50
'thickness_scintillator_1': random.uniform(0.1, 50),
'thickness_scintillator_2': random.uniform(0.1, 50),
'material_abs_0': random.uniform(-1, 1), # Random float between 0.1 and 1
'material_abs_1': random.uniform(-1, 1),
'material_abs_2': random.uniform(-1, 1),
'material_scint_0': -0.3,#random.uniform(-1, 1), # Random float between 0.1 and 1
'material_scint_1': random.uniform(-1, 1),
'material_scint_2': random.uniform(-1, 1),
}
# start a new wandb run to track this script
wandb.init(
# set the wandb project where this run will be logged
project="Culprit",
name = 'Run' + str(run_num ), #Changed
# track hyperparameters and run metadata
config={
'start_pars' : start_pars,
'seed' : (args.seed)
}
)
box = np.array(len(start_pars)*[1.5]) #will be adjusted to some extent by optimizer
#start parameters from run 8
#start_pars = {'thickness_absorber_0': 0.7642903, 'thickness_absorber_1': 10.469371, 'thickness_scintillator_0': 30.585306, 'thickness_scintillator_1': 22.256506}
gen = Generator(box,
Parameters(parameters = start_pars),
Parameters(ToP, parameters = start_pars),
types = ToP,
n_vars = 40,
n_events_per_var = 400//divide,
particles=[['gamma',0.22,'pi+',2.11]])
particles=['gamma',0.22,'pi+',2.11])
# reconstruction
reco_model = Reconstruction(gen.n_parameters,
gen.n_parameters_per_sensor * gen.n_sensors,
......@@ -86,15 +112,13 @@ if __name__ == "__main__":
n_time_steps=50)
optimizer = Optimizer(gen, surrogate_model, reco_model, gen.parameters,
constraints= {'length': 200, 'upper' : 1, 'lower': 0,'diff' : 0.66, 'cost' : 50000}) # length of detector, upper limit, lower limit and max diff of material parameters
constraints= {'length': 200, 'upper' : 1, 'lower': -1,'cost' : 50000}) # length of detector, upper limit, lower limit and max diff of material parameters
import sys
n_epochs_pre = 25//divide
n_epochs_pre = 24//divide
n_epochs_main = 40//divide
parameters = gen.parameters
parameters = gen.parameters
#if saved evolution exists, load it
import os.path
......@@ -110,6 +134,8 @@ if __name__ == "__main__":
save it in human readable format as well
'''
ps_dict = gen.translate_parameters(parameters)
wandb.log(ps_dict)
ps = {k: ps_dict[k] for k in ps_dict.keys()}
ps.update(
{'reco_model': reco_model.state_dict(),
......@@ -151,6 +177,7 @@ if __name__ == "__main__":
plt.legend()
plt.xlabel('step')
plt.savefig(outpath+'evolution_continuous.png')
wandb.log({"evolution_continuous": wandb.Image(plt)})
plt.close()
for k in pltdict.keys():
......@@ -165,9 +192,6 @@ if __name__ == "__main__":
plt.savefig(outpath+'evolution_material.png')
plt.close()
#save_evolution(parameters, 1., 1.)
#exit()
def make_check_plot(surr_out, reco_out, true_in, evolution, outname="" ):
'''
......@@ -175,23 +199,63 @@ if __name__ == "__main__":
'''
reco_resp = (reco_out.detach().cpu().numpy() - true_in.detach().cpu().numpy()) / true_in.detach().cpu().numpy()
surr_resp = (surr_out.detach().cpu().numpy() - true_in.detach().cpu().numpy()) / true_in.detach().cpu().numpy()
#range = np.min([np.min(reco_resp), np.min(surr_resp)]), np.max([np.max(reco_resp), np.max(surr_resp)])
if len(outname) == 0:
outname = str(len(evolution))
# if len(outname) == 0:
# outname = str(len(evolution))
# Compute histograms
reco_hist, bins = np.histogram(reco_resp.flatten(), bins=51, range=(-1, 1))
surr_hist, _ = np.histogram(surr_resp.flatten(), bins=51, range=(-1, 1))
plt.figure(figsize=(10,10))
plt.hist(reco_resp.flatten(), bins=51, range=(-1,1), alpha = 0.5, label='G4+reco')
plt.hist(surr_resp.flatten(), bins=51, range=(-1,1), alpha = 0.5, label='E2E surrogate')
plt.text(0.15, 0.8, 'step '+outname, transform=plt.gca().transAxes)
if len(evolution):
plt.text(0.15, 0.7, 'o_loss '+str((evolution[-1][1])), transform=plt.gca().transAxes)
plt.legend()
plt.savefig(outpath+'response_'+outname+'_lin.png')
plt.yscale('log')
plt.savefig(outpath+'response_'+outname+'_log.png')
plt.close()
# Load existing data if the JSON file exists
pickle_path = os.path.join(outpath, 'responses.pkl')
# Load existing data or create empty list
if os.path.exists(pickle_path):
with open(pickle_path, 'rb') as f:
all_histogram_data = pickle.load(f)
else:
all_histogram_data = []
# if(len(evolution)==0):
# return
# Append new histogram data
histogram_data = {
"iteration": len(evolution),
"reco_hist": reco_hist.tolist(),
"surr_hist": surr_hist.tolist(),
"bins": bins.tolist(),
"evolution_loss": evolution[-1][1] if evolution else None
}
all_histogram_data.append(histogram_data)
# Save the combined histogram data
with open(pickle_path, 'wb') as f:
pickle.dump(all_histogram_data, f)
with open(outpath+'responses.txt', 'w') as f:
for e in all_histogram_data:
f.write('{')
for key, value in e.items():
# Format key-value pair into a string and write to file
if (key != "evolution_loss"):
f.write(f"{key}: {value},")
else:
f.write(f"{key}: {value}" + "}")
f.write('\n')
# plt.figure(figsize=(10,10))
# plt.hist(reco_resp.flatten(), bins=51, range=(-1,1), alpha = 0.5, label='G4+reco')
# plt.hist(surr_resp.flatten(), bins=51, range=(-1,1), alpha = 0.5, label='E2E surrogate')
# plt.text(0.15, 0.8, 'step '+outname, transform=plt.gca().transAxes)
# if len(evolution):
# plt.text(0.15, 0.7, 'o_loss '+str((evolution[-1][1])), transform=plt.gca().transAxes)
# plt.legend()
# plt.savefig(outpath+'response_'+outname+'_lin.png')
# plt.yscale('log')
# plt.savefig(outpath+'response_'+outname+'_log.png')
# plt.close()
#pre-train reco and surrogate model
def pre_train():
......@@ -237,10 +301,15 @@ if __name__ == "__main__":
return loss < 4. * best_surrogate_loss
covariance_pcas = []
results = []
scale_update = 0
for i in range(1000): #just an example
ds = gen.create_torch_dataset(parameters) # full variations plus a tiny bit
model_means = ds.c_means
model_stds = ds.c_stds
reco_model.to('cuda')
reco_model.train_model(ds, batch_size=256, n_epochs= n_epochs_main//4, lr=0.003)
reco_model.train_model(ds, batch_size=1024, n_epochs= n_epochs_main//2, lr=0.001)
......@@ -269,7 +338,8 @@ if __name__ == "__main__":
print("Optimization")
#apply the surrogate model
updated_detector_parameters, optimal, o_loss = optimizer.optimize(surrogate_dataset, batch_size=512, n_epochs=40, lr=0.03, add_constraints=True)#some noise
updated_detector_parameters, optimal, o_loss, reco_surrogate_loss, constraint_loss, scale_update = optimizer.optimize(surrogate_dataset,ToP, scale_update, batch_size=512, n_epochs=40, lr=0.02, add_constraints=True)#some noise
print("Updated detector parameters: ", gen.translate_parameters(updated_detector_parameters), "; not interrupted", optimal, "; change", updated_detector_parameters - parameters)
......@@ -286,10 +356,74 @@ if __name__ == "__main__":
pre_train()
sl = 0. #reset surrogate loss too to avoid infinite loop
continue
#
model_means = [tensor.to('cpu').numpy() for tensor in model_means]
model_stds = [tensor.to('cpu').numpy() for tensor in model_stds]
single_ds = gen.create_torch_dataset(updated_detector_parameters, random_scaler=0, means = model_means, stds = model_stds, single=True)
y_preds, loss = reco_model.apply_model_in_batches(single_ds, 40)
detector_array = single_ds.detector_array
context_array = single_ds.context_array
true_energy = single_ds.target_array
# print('Detector Shape is : ', detector_array.shape ,'\n' , context_array.shape , '\n' , true_energy.shape)
context_array = context_array.reshape(1,-1)
true_energy = true_energy.reshape(1,-1)
y_preds_cpu = y_preds.detach().cpu()
# # Convert to numpy arrays
reco_energy = y_preds_cpu.numpy()
reco_energy = reco_energy.reshape(1, -1)
iteration_dict = {
'detector_array': detector_array,
'context_array': context_array,
'true_energy': true_energy,
'reco_energy': reco_energy,
'loss': loss,
'reco_loss': reco_loss,
'o_loss' : o_loss,
'reco_surrogate_loss' : reco_surrogate_loss,
'constraint_loss': constraint_loss,
'scale_update' : scale_update
}
scale = 1+ (0.02*scale_update)
wandb_dict = {'loss': loss,'o_loss': o_loss,'reco_loss' : reco_loss ,'reco_surrogate_loss' : reco_surrogate_loss,
'constraint_loss': constraint_loss, 'scale': scale}
wandb.log(wandb_dict)
# Append the dictionary to the results list
results.append(iteration_dict)
pickle_file = 'results.pkl'
with open(outpath + pickle_file, 'wb') as f:
pickle.dump(results, f)
# step more in the direction.
parameters = updated_detector_parameters + 0.5*(updated_detector_parameters - parameters)
txt_file = 'results.txt'
# Open the file in write mode
with open(outpath + txt_file, 'w') as f:
for i, result in enumerate(results):
f.write(f"Iteration {i + 1}:\n")
f.write(f"Detector Array:\n{result['detector_array']}\n")
f.write(f"Context Array:\n{result['context_array']}\n")
f.write(f"True Energy:\n{result['true_energy']}\n")
f.write(f"Reco Energy:\n{result['reco_energy']}\n")
f.write(f"Loss:\n{result['loss']}\n")
f.write(f"Optimizer Loss:\n{result['o_loss']}\n")
f.write(f"Scale Update:\n{result['scale_update']}\n")
f.write("\n")
print(f"Saved results to {txt_file}")
parameters = updated_detector_parameters + 0.5*(updated_detector_parameters - parameters)
# for bookkeeping, append the generator box_covariance to a text file
with open(outpath+'box_covariance.txt', 'a') as f:
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment