def walker(*args):
dirname, dirnames, filenames = args[0][0], args[0][1], args[0][2]
clsname = os.path.split(dirname)
if len(filenames) < 1:
return 0
phsp = PhaseSpace(clsname[-1])
for filename in filenames:
if 'txt' in filename:
filepath = os.path.join(dirname, filename)
with open(filepath, 'r') as file:
file.readline(), file.readline(), file.readline()
phsp.append_data(file.readlines())
rph.plotclsphsp(phsp, output)
rph.plotclshistogram(phsp, output)
def distribution_to_sixtrack(dist_path, save_file):
ps = PhaseSpace(dist_path)
with open(save_file, 'w') as f:
for i in range(ps.nbr_p):
x, px, y, py = 0, 0, 0, 0
z = phi_to_z(ps.phase[i])
E = 450e3 + ps.denergy[i] * 1e-6 # MeV
f.write(
" {:>17.12f}, {:>17.12f}, {:>17.12f}, {:>17.12f}, {:>17.12f}, {:>17.12f}\n"
.format(x, px, y, py, z, E))
def aggregate(self):
n = 1
required_files = [
"stdout.txt", settings.STARTDIST_FILE, settings.COLL_FILE,
"TCPc_IR7.ch"
]
lg.log("aggregating batch from '{}'".format(self.path))
while os.path.isdir("{}/{}{}".format(self.path,
settings.BATCH_JOB_PRESTRING, n)):
job_path = "{}/{}{}".format(self.path,
settings.BATCH_JOB_PRESTRING, n)
lg.log("trying '{}'...".format(job_path), end=" ")
self.nbr_jobs += 1
for req in required_files:
if not os.path.exists("{}/{}".format(job_path, req)):
lg.log("not ok",
module_prestring=False,
log_level=LogLevel.warning)
break
else:
pid_offset = self.ps.nbr_p if self.ps else 0
ps = PhaseSpace("{}/{}".format(job_path,
settings.STARTDIST_FILE),
mute=True)
hm = lm.CHitMap("{}/{}".format(job_path, settings.COLL_FILE),
pid_offset=pid_offset,
store_ip=False)
if self.ps is None:
self.ps = ps
self.hitmap = hm
else:
self.ps = PhaseSpace.merge_two(self.ps, ps)
self.hitmap = self.hitmap.concatenate(hm)
self.nbr_valid_jobs += 1
lg.log("ok",
module_prestring=False,
log_level=LogLevel.success)
n += 1
def phasespace_evolution():
frames = 300
ps = PhaseSpace(None)
ps.create_plot()
ps.format_axes()
ps.plot_trajectory("{}/0lines.dat".format(phasespace_dir))
ani = animation.FuncAnimation(ps.fig,
phasespace_frame,
int(frames),
fargs=(ps, ),
interval=100,
blit=False)
mov_file = "{}/mov.mp4".format(phasespace_dir)
lg.log("saving movie to '{}'".format(mov_file))
ani.save(mov_file)
lg.log("finished saving to '{}'".format(mov_file))
ax3.plot(H, list(map(curve, H)), color="darkorange")
ax4.plot(H, list(map(curve, H)), color="darkorange")
else:
raise Exception("plot_type not valid")
if block:
plt.show()
else:
plt.draw()
if __name__ == "__main__":
action = sys.argv[1]
fill = af.aggregate_fill(1, from_cache=True)
ps = PhaseSpace(settings.STARTDIST_PATH)
hitmap = lm.CHitMap(settings.COLL_PATH)
comp = LHCComparison(fill, ps, hitmap)
if action == "fit":
blm_s = (comp.t(), comp.BLM())
comp.set_window(14, 20)
# comp.set_window(0, 14)
# comp.halign()
comp.fit_action_values()
r = comp.fit_results()
coef = r['c']
avs = r['action_values']
with open(FIT_COEF_FILE, "w") as f:
f.write("# {:>3} {:<20} {}\n".format("n", "Coefficients", "∆H"))
blit=False)
mov_file = "{}/mov.mp4".format(phasespace_dir)
lg.log("saving movie to '{}'".format(mov_file))
ani.save(mov_file)
lg.log("finished saving to '{}'".format(mov_file))
SAVE_FILE = ''
if (len(argv) > 2):
SAVE_FILE = argv[2]
if __name__ == "__main__":
ACTION = argv[1]
if ACTION == "animate":
lg.log("animate trajectory")
ps = PhaseSpace(settings.PARTICLE_PATH)
ps.animate(get_lossmap(settings.COLL_PATH), save_to=SAVE_FILE)
if ACTION == "animate-full":
lg.log("animating trajectory and background")
output = "calc/animate-full.mp4"
lg.log("will save movie to '{}'".format(output))
ps = PhaseSpace(settings.PARTICLE_PATH)
ps.animate(get_lossmap(settings.COLL_PATH),
output,
animateBackground=True)
elif ACTION == "lossmap":
lg.log("plot lossmap")
lossmap = get_lossmap(settings.COLL_PATH)
plot_lossmap([lossmap], ["Toy model"], SAVE_FILE)
elif ACTION == "lossmap2":
lg.log("plot lossmap2")
import matplotlib.pyplot as plt
import numpy as np
from scipy import interpolate
from scipy.optimize import least_squares
from settings import settings
dir_in = "/Users/swretbor/Workspace/collimation/proj/synchrotron/simulations/store/pdf_inside"
dir_out = "/Users/swretbor/Workspace/collimation/proj/synchrotron/simulations/store/pdf_outside"
# dir_in = "/Users/swretbor/Workspace/collimation/proj/synchrotron/simulations/store/pdf_inside_1turnmap"
# dir_out = "/Users/swretbor/Workspace/collimation/proj/synchrotron/simulations/store/pdf_outside_1turnmap"
fill = af.aggregate_fill(1, from_cache=True)
inside = {
"ps": PhaseSpace(dir_in + "/" + settings.STARTDIST_FILE),
"lm": lm.CHitMap(dir_in + "/" + settings.COLL_FILE)
}
ic = LHCComparison(fill, inside['ps'], inside['lm'])
outside = {
"ps": PhaseSpace(dir_out + "/" + settings.STARTDIST_FILE),
"lm": lm.CHitMap(dir_out + "/" + settings.COLL_FILE)
}
oc = LHCComparison(fill, outside['ps'], outside['lm'])
# ic.halign()
# oc.halign()
tmin = min(ic.t().min(), oc.t().min())
tmax = max(ic.t().max(), oc.t().max())
labels = [
# "CInside_C",
"Constant",
# "CInside_LD",
"Lin. decrease",
# "CInside_LI",
"Lin. increase"
]
pss = []
hms = []
for d in dirs:
print("Reading '{}'".format(d))
hm = lm.CHitMap(d + "/" + settings.COLL_FILE)
hms.append(hm)
pss.append(PhaseSpace(d + "/" + settings.STARTDIST_FILE))
print("Plot")
fill = af.aggregate_fill(1, from_cache=True)
fig, loss_ax = plt.subplots()
loss_ax.plot(*fill.blm_ir3(),
label="Aggr. fill (beam {})".format(fill.beam),
color='r')
loss_ax.set_yscale("log")
loss_ax.set_xlim([-5, 40])
loss_ax.set_ylim([0.5e-5, 1])
loss_ax.set_ylabel("Losses (∆particles/1.3s)")
loss_ax.set_xlabel("t (s)")
loss_ax.axvspan(0.0,