def plot_item(self, y_name, get_item_lambda, axis, canvas, mgraph):
x_data = [item[self.subs_key][axis] for item in self.data]
x_name = utilities.sub_to_name(axis)+utilities.sub_to_units(axis)
y_data = []
for item in self.data:
try:
y_data.append(get_item_lambda(item))
except Exception:
sys.excepthook(*sys.exc_info())
y_data.append(0.)
name = utilities.sub_to_name(axis)+" vs "+y_name
print(" ", y_name, y_data)
hist, graph = xboa.common.make_root_graph(name, x_data, x_name, y_data, y_name)
if canvas == None:
canvas = xboa.common.make_root_canvas(name)
if self.log_x:
canvas.SetLogx(True)
canvas.Draw()
mgraph = ROOT.TMultiGraph()
mgraph.SetTitle(";"+x_name+";"+y_name)
self.root_objects.append(mgraph)
if not mgraph.GetListOfGraphs():
index = 0
else:
index = mgraph.GetListOfGraphs().GetEntries()
color_list = [ROOT.kRed, ROOT.kBlue]
if len(x_data) < 25:
style_list = [20, 20]
else:
style_list = [7, 7]
graph.SetMarkerColor(color_list[index])
graph.SetMarkerStyle(style_list[index])
mgraph.Add(graph)
return canvas, mgraph
def track_one(self, index, subs_overrides_list, kinetic_energy=None):
if kinetic_energy == None:
kinetic_energy = self.e0
subs = self.config.substitution_list[index]
for sub_dict in subs_overrides_list:
for item, key in sub_dict.items():
subs[item] = key
print("Tracking with", end=' ')
for key in sorted(subs.keys()):
print(utilities.sub_to_name(key), subs[key], end=' ')
print()
self._temp_dir() # changes to tmp/find_rf_parameters/
xboa.common.substitute(self.config.tracking["lattice_file"],
"SectorFFAGMagnet.tmp", subs)
ref_hit = sorted(
self.closed_orbit_list[index],
key=lambda hit: abs(hit["kinetic_energy"] - self.e0))[0]
test_hit = sorted(
self.closed_orbit_list[index],
key=lambda hit: abs(hit["kinetic_energy"] - kinetic_energy))[0]
ref_hit = ref_hit.deepcopy()
test_hit = test_hit.deepcopy()
momentum = ref_hit["p"]
ref_hit["x"] = 0.
ref_hit["px"] = 0.
ref_hit["p"] = momentum
ref_hit.mass_shell_condition("p")
tracking = OpalTracking("SectorFFAGMagnet.tmp", "disttest.dat",
ref_hit,
self.config.find_rf_parameters["probe_files"],
self.config.tracking["opal_path"], "log")
print("Tracking hit with kinetic energy:", test_hit["kinetic_energy"])
hit_list = tracking.track_one(test_hit)
return hit_list
def plot_data_2d(data, tune_axis, plot_dir):
x_tune = [item['x_tune'] for item in data]
y_tune = [item['y_tune'] for item in data]
print("X Tunes", x_tune)
print("Y Tunes", y_tune)
x_name = "Varying "
axis_candidates = utilities.get_substitutions_axis(data)
for key in axis_candidates:
min_val = str(min(axis_candidates[key]))
max_val = str(max(axis_candidates[key]))
x_name += utilities.sub_to_name(
key) + " from " + min_val + " to " + max_val
axis_candidates = utilities.get_substitutions_axis(data)
canvas = xboa.common.make_root_canvas('tune x vs y')
hist, graph = xboa.common.make_root_graph('tune x vs y',
x_tune,
"horizontal tune",
y_tune,
"vertical tune",
xmin=0.,
xmax=1.,
ymin=0.,
ymax=1.)
hist.SetTitle(x_name)
hist.Draw()
#utilities.tune_lines(canvas)
graph.SetMarkerStyle(24)
graph.Draw("P SAME")
canvas.Update()
for format in "eps", "png", "root":
canvas.Print(plot_dir + "/tune_x_vs_y." + format)
def run_tracking(self, index):
opal_exe = self.config.tracking["opal_path"]
input_file = self.config.tracking["lattice_file"]
n_cores = self.config.tracking["n_cores"]
mpi_exe = self.config.tracking["mpi_exe"]
lattice_file = self.run_dir+'SectorFFAGMagnet.tmp'
subs = self.config.substitution_list[index]
for key, value in self.config.track_beam["subs_overrides"].items():
subs[key] = value
xboa.common.substitute(input_file, lattice_file, subs)
log_name = self.run_dir+"/log"
ref_hit = self.reference()
probe_files = self.config.track_beam["probe_files"]
self.tracking = OpalTracking("SectorFFAGMagnet.tmp", 'disttest.dat', ref_hit, probe_files, opal_exe, log_name, None, n_cores, mpi_exe)
tunes_analysis = TunesAnalysis(self.config)
phase_space_plots = PhaseSpacePlots(self.config)
tunes_analysis.set_match(self.centre[0:4], self.ellipse[0:4, 0:4])
if self.config.track_beam["do_track"]:
print("Running tracking with\n ", end=' ')
for key, value in subs.items():
print(utilities.sub_to_name(key)+":", value, end=' ')
print()
self.tracking.track_many(self.hits_in, None)
print(os.getcwd(), probe_files)
self.tracking._read_probes(tunes_analysis)
def setup_subs(self, fields):
subs = self.config.substitution_list[0]
if self.first_tracking:
for key in sorted(subs.keys()):
print(utilities.sub_to_name(key), subs[key], end=' ')
self.first_tracking = False
self.overrides.update(fields)
print("OVERRIDES", self.overrides)
utilities.do_lattice(self.config, self.subs, self.overrides)
def make_global_plot(self, min_n_points, plot_dir):
axis_candidates = utilities.get_substitutions_axis(
self.data, "substitutions")
my_axes_x = dict([(key, []) for key in axis_candidates])
da = {'y_da': [], 'x_da': []}
for item in self.data:
for da_key in da:
if da_key not in list(item.keys()):
print("Did not find", da_key, "in keys", list(item.keys()),
"... skipping")
continue
da[da_key].append(
self.calculate_da(item, da_key, min_n_points) * 1e3)
print(da_key, da[da_key][-1])
for ax_key in my_axes_x:
my_axes_x[ax_key].append(item['substitutions'][ax_key])
for ax_key in my_axes_x:
ordinate = my_axes_x[ax_key]
x_name = utilities.sub_to_name(ax_key)
canvas_name = x_name + ' vs da'
canvas = common.make_root_canvas(canvas_name)
canvas.SetLogy()
x_range = [min(ordinate), max(ordinate)]
y_range = [
min(da['x_da'] + da['y_da']),
max(da['x_da'] + da['y_da'])
]
y_range[0] = max(y_range[0] / 2, 0.01)
hist, graph = common.make_root_graph('axes',
x_range,
x_name,
y_range,
"DA [#pi #mum]",
ymin=y_range[0],
ymax=y_range[1] * 3)
hist.Draw()
if len(da['x_da']) == len(ordinate):
hist_x, graph_x = common.make_root_graph(
'horizontal da', ordinate, "", da['x_da'], "")
graph_x.SetMarkerStyle(24)
graph_x.SetMarkerColor(2)
graph_x.Draw("SAME P")
if len(da['y_da']) == len(ordinate):
hist_y, graph_y = common.make_root_graph(
'vertical da', ordinate, "", da['y_da'], "")
graph_y.SetMarkerStyle(26)
graph_y.SetMarkerColor(4)
graph_y.Draw("SAME P")
#legend = common.make_root_legend(canvas, [graph_x, graph_y])
#legend.Draw()
canvas.Update()
canvas_name = canvas_name.replace(" ", "_")
for format in "eps", "png", "root":
canvas.Print(plot_dir + "/" + canvas_name + "." + format)
return
def find_tune_dphi(self):
"""
Algorithm is to just calculate the turn-by-turn phase advance; this is
done by evolving turn-by-turn the track; calculating a matched ellipse
by looking at tracking output; transforming the ellipse into a circle
using LU decomposition; then calculating the angle advanced.
"""
cwd = os.getcwd()
fout = open(self.output_filename, "w")
index = 0
for i, closed_orbit in enumerate(self.closed_orbits_cached):
if self.row != None and i not in self.row:
continue
if len(closed_orbit["hits"]) == 0:
print("Error - no closed orbit")
continue
index += 1
if index >= self.config.find_tune["root_batch"]:
ROOT.gROOT.SetBatch(True)
subs = closed_orbit["substitutions"]
for item, key in self.config.find_tune["subs_overrides"].items():
subs[item] = key
print("Finding optics with", end=' ')
for key in sorted(subs.keys()):
print(utilities.sub_to_name(key), subs[key], end=' ')
print()
tune_info = {"substitutions": subs}
for axis1, axis2, delta1, delta2 in [("x", "px", self.delta_x, 0.),
("y", "py", self.delta_y, 0.)
]:
if self.do_axis != None and axis1 != self.do_axis:
continue
hit = Hit.new_from_dict(closed_orbit["hits"][0])
self._temp_dir()
common.substitute(self.lattice_src,
self.tmp_dir + self.lattice, subs)
tracking = OpalTracking(self.tmp_dir + self.lattice,
self.tmp_dir + self.beam_file,
self._reference(hit),
self.probe_filename, self.opal,
self.tmp_dir + self.log_file)
tracking.clear_path = self.tmp_dir + "/*.loss"
for key in sorted(tune_info.keys()):
if "signal" not in key and "dphi" not in key:
print(" ", key, tune_info[key])
print(json.dumps(tune_info), file=fout)
fout.flush()
os.chdir(cwd)
def setup_subs(self, fields):
subs = self.config.substitution_list[0]
if self.first_tracking:
for key in sorted(subs.keys()):
print(utilities.sub_to_name(key), subs[key], end=' ')
self.first_tracking = False
for key, value in self.overrides.items():
self.subs[key] = value
for i, a_field in enumerate(fields):
self.subs["__bump_field_" + str(i + 1) + "__"] = a_field
self.subs["__phi_foil_probe__"] = \
self.config.find_bump_parameters["foil_probe_phi"]
xboa.common.substitute(self.config.tracking["lattice_file"],
self.tmp_dir + 'SectorFFAGMagnet.tmp',
self.subs)
def setup_tracking(self, co_element):
subs = co_element["substitutions"]
for item, key in self.config.find_da["subs_overrides"].items():
subs[item] = key
print("Set up tracking for da with", end=' ')
for key in sorted(subs.keys()):
print(utilities.sub_to_name(key), subs[key], end=' ')
self.ref_hit = self.reference(co_element["hits"][0])
lattice_src = self.config.tracking["lattice_file"]
common.substitute(lattice_src, self.run_dir + "/SectorFFAGMagnet.tmp",
subs)
tracking_file = self.config.find_da["probe_files"]
self.tracking = OpalTracking(self.run_dir + "/SectorFFAGMagnet.tmp",
self.tmp_dir + '/disttest.dat',
self.ref_hit, tracking_file,
self.opal_exe, self.tmp_dir + "/log")
def plot_tune(data, tune_axis, plot_dir):
print("Plotting tune")
axis_candidates = utilities.get_substitutions_axis(data)
metres = 1e-3
for key in axis_candidates:
sys.stdout.flush()
x_name = utilities.sub_to_name(key)
x_list = axis_candidates[key]
y_list = [filtered_dphi(item[tune_axis]) for item in data]
x_name += utilities.sub_to_units(key)
fig_index = matplotlib_wrapper.make_graph(x_list, x_name, y_list,
tune_axis)
for format in ["png"]:
name = plot_dir + "/" + tune_axis + "." + format
fig = matplotlib.pyplot.figure(fig_index)
fig.savefig(name)
def do_plots_by_sub(self):
subs_axes = utilities.get_substitutions_axis(self.data, self.subs_key)
print(subs_axes)
for axis in subs_axes:
axis_name = utilities.sub_to_name(axis).lower().replace(" ", "_")
for name in "x", "x'", "y", "y'":
canvas, mgraph = self.plot_item(
name,
lambda item: Hit.new_from_dict(item["seed_hit"])[name],
axis, None, None)
mgraph.Draw("AP")
name = name.replace("'", "p")
canvas.Print(self.output_dir+"/"+axis_name+"_vs_"+name+".png")
canvas, mgraph = None, None
for self.phase_advance_axis in 0, 1:
self.tm_lambda = self.tm_phase_advance
name = "Phase Advance "+str(self.phase_advance_axis)
canvas, mgraph = self.plot_item(name, self.get_tm, axis, canvas, mgraph)
mgraph.Draw("AP")
canvas.Print(self.output_dir+"/"+axis_name+"_vs_phase_advance.png")
for self.phase_advance_axis in 0, 1:
canvas, mgraph = None, None
self.tm_lambda = self.tm_phase_advance
name = "phase advance "+str(self.phase_advance_axis)
canvas, mgraph = self.plot_item(name, self.get_tm, axis, canvas, mgraph)
mgraph.Draw("AP")
name = name.replace(" ", "_")
canvas.Print(self.output_dir+"/"+axis_name+"_vs_"+name+".png")
self.tm_lambda = self.tm_det
name = "Determinant"
canvas, mgraph = self.plot_item(name, self.get_tm, axis, None, None)
mgraph.Draw("AP")
canvas.Print(self.output_dir+"/"+axis_name+"_vs_determinant.png")
name = "delta"
canvas, mgraph = self.plot_item(name, self.get_co, axis, None, None)
mgraph.Draw("AP")
#canvas.SetLogy()
canvas.Print(self.output_dir+"/"+axis_name+"_vs_delta.png")
def get_groups(data, group_axis):
axis_candidates = utilities.get_substitutions_axis(data)
if group_axis != None and group_axis not in axis_candidates:
raise RuntimeError("Did not recognise group axis " + str(group_axis))
if group_axis == None:
group_list = [item for item in data]
return {"": {"item_list": [i for i in range(len(data))]}}
else:
# group_list is lists the possible groups
# e.g. list of all possible values of "__bump_field_1__"
group_list = [item['substitutions'][group_axis] for item in data]
group_list = list(set(group_list)) # unique list
# tmp_group_dict is mapping from group value to the items having that value
tmp_group_dict = dict([(group, []) for group in group_list])
for key in tmp_group_dict:
tmp_group_dict[key] = [i for i, item in enumerate(data) \
if item['substitutions'][group_axis] == key]
group_dict = {}
for key in tmp_group_dict:
new_key = utilities.sub_to_name(group_axis) + " " + format(key, "3.3g")
group_dict[new_key] = {'item_list': tmp_group_dict[key]}
print(new_key, ":", group_dict[new_key])
return group_dict
def get_axes(axis_key, group_dict, tune_axis):
x_name = utilities.sub_to_name(axis_key)
canvas_name = tune_axis + ' vs ' + x_name
canvas = xboa.common.make_root_canvas(canvas_name)
tune = []
ordinate = []
if x_name in units:
x_name += " [" + units[x_name] + "]"
for group_key in group_dict:
tune += group_dict[group_key]['x_tune'] + group_dict[group_key][
'y_tune']
ordinate += group_dict[group_key]['axes'][axis_key]
y_min, y_max = min(tune), max(tune)
x_min, x_max = min(ordinate), max(ordinate)
if x_min == x_max:
x_min -= 0.5
x_max += 0.5
if y_min == y_max:
y_min -= 0.5
y_max += 0.5
dx, dy = x_max - x_min, y_max - y_min
x_min -= dx * 0.1
x_max += dx * 0.4
y_min -= dy * 0.1
y_max += dy * 0.1
hist = xboa.common.make_root_histogram('axes', [x_min - 1],
x_name,
1000, [y_min - 1],
tune_axis,
1000,
xmin=x_min,
xmax=x_max,
ymin=y_min,
ymax=y_max)
hist.Draw()
return canvas, hist
def find_tune_dphi(self):
"""
Algorithm is to just calculate the turn-by-turn phase advance; this is
done by evolving turn-by-turn the track; calculating a matched ellipse
by looking at tracking output; transforming the ellipse into a circle
using LU decomposition; then calculating the angle advanced.
"""
cwd = os.getcwd()
fout = open(self.output_filename, "w")
index = 0
for i, closed_orbit in enumerate(self.closed_orbits_cached):
if self.row != None and i not in self.row:
continue
if len(closed_orbit["hits"]) == 0:
print("Error - no closed orbit")
continue
index += 1
subs = closed_orbit["substitutions"]
for item, key in self.config.find_tune["subs_overrides"].items():
subs[item] = key
print("Finding tune with", end=' ')
for key in sorted(subs.keys()):
print(utilities.sub_to_name(key), subs[key], end=' ')
print()
tune_info = {"substitutions": subs}
for axis1, axis2, delta1, delta2 in [("x", "px", self.delta_x, 0.),
("y", "py", self.delta_y, 0.)
]:
if self.do_axis != None and axis1 != self.do_axis:
continue
hit = Hit.new_from_dict(closed_orbit["hits"][0])
self._temp_dir()
common.substitute(self.lattice_src,
self.tmp_dir + self.lattice, subs)
tracking = OpalTracking(self.tmp_dir + self.lattice,
self.tmp_dir + self.beam_file,
self._reference(hit),
self.probe_filename, self.opal,
self.tmp_dir + self.log_file)
tracking.clear_path = self.tmp_dir + "/*.loss"
finder = DPhiTuneFinder()
try:
finder.run_tracking(axis1, axis2, delta1, delta2, hit,
tracking)
except RuntimeError:
sys.excepthook(*sys.exc_info())
for track_index, track in enumerate(tracking.last):
print('Track', track_index, 'of', len(tracking.last), \
'with', len(track), 'hits')
finder.u = finder.u[1:]
finder.up = finder.up[1:]
try:
tune = finder.get_tune(subs["__n_turns__"] / 10.)
except:
tune = 0.
print(' Found', len(finder.dphi), 'dphi elements with tune',
tune, "+/-", finder.tune_error)
tune_info[axis1 + "_tune"] = tune
tune_info[axis1 + "_tune_rms"] = finder.tune_error
tune_info[axis1 + "_signal"] = list(zip(finder.u, finder.up))
tune_info[axis1 + "_dphi"] = finder.dphi
tune_info[axis1 + "_n_cells"] = len(finder.dphi)
self.do_plots(i, axis1, axis2, finder)
for i, u in enumerate([]): #finder.u[:-1]):
up = finder.up[i]
dphi = finder.dphi[i]
t = finder.t[i]
u_chol = finder.point_circles[i][0]
up_chol = finder.point_circles[i][1]
phi = math.atan2(up_chol, u_chol)
print(str(i).ljust(4), str(round(t, 4)).rjust(8), "...", \
str(round(u, 4)).rjust(8), str(round(up, 4)).rjust(8), "...", \
str(round(u_chol, 4)).rjust(8), str(round(up_chol, 4)).rjust(8), "...", \
str(round(phi, 4)).rjust(8), str(round(dphi, 4)).rjust(8))
for key in sorted(tune_info.keys()):
if "signal" not in key and "dphi" not in key:
print(" ", key, tune_info[key])
print(json.dumps(tune_info), file=fout)
fout.flush()
os.chdir(cwd)
def find_closed_orbit(self, sub_index, subs, seed):
"""
Find the closed orbit
- sub_index: indexes element in the substitution loop (for scans)
- subs: dictionary of key value pairs for substitution into the lattice
- seed: best guess of the closed orbit
Returns the trajectory of the closed orbit
"""
max_iterations = self.config.find_closed_orbits["max_iterations"]
probe = self.config.find_closed_orbits["probe_files"]
for key in sorted(subs.keys()):
print(utilities.sub_to_name(key), subs[key], end=' ')
print()
utilities.clear_dir(self.run_dir)
os.chdir(self.run_dir)
print("Running in", os.getcwd())
common.substitute(self.config.tracking["lattice_file"],
self.config.tracking["lattice_file_out"], subs)
energy = subs["__energy__"]
ref_hit = utilities.reference(self.config, energy)
tracking = utilities.setup_tracking(self.config, probe, energy)
seed_hit = ref_hit.deepcopy()
seed_hit["x"] = seed[0]
seed_hit["px"] = seed[1]
# fix momentum
seed_hit["pz"] = (ref_hit["p"]**2 - seed_hit["px"]**2)**0.5
print("Reference kinetic energy:", ref_hit["kinetic_energy"])
print("Seed kinetic energy: ", seed_hit["kinetic_energy"])
finder = EllipseClosedOrbitFinder(tracking, seed_hit)
generator = finder.find_closed_orbit_generator(["x", "px"], 1)
x_std_old = 1e9
i = -1
will_loop = True
iteration = None
while will_loop:
try:
iteration = next(generator)
except (StopIteration, RuntimeError):
will_loop = False
print(sys.exc_info()[1])
i += 1
heading = [
'station', 't', 'x', 'px', 'y', 'py', 'z', 'pz', 'r', 'pt',
'kinetic_energy'
]
for key in heading:
print(str(key).rjust(10), end=' ')
print()
for hit in tracking.last[0]:
for key in heading:
print(str(round(hit[key], 1)).rjust(10), end=' ')
print()
if iteration == None:
continue
print(iteration.centre)
#if iteration.centre != None: #i == 0 and
if i == 0:
self.plot_iteration(sub_index, i, iteration, energy)
if i >= max_iterations:
break
x_mean = numpy.mean([point[0] for point in iteration.points])
x_std = numpy.std([point[0] for point in iteration.points])
print("Seed:", iteration.points[0][0], "Mean:", x_mean, "Std:",
x_std)
if type(iteration.centre) != type(
None) and x_std >= x_std_old: # require convergence
break
x_std_old = x_std
os.chdir(self.out_dir)
if i > 0:
self.plot_iteration(sub_index, i, iteration, energy)
return tracking.last[0]
def get_title(variables):
title = ""
for key in variables:
title += utilities.sub_to_name(key) + ": " + str(variables[key])
return title
def find_tune_dphi(self):
"""
Algorithm is to just calculate the turn-by-turn phase advance; this is
done by evolving turn-by-turn the track; calculating a matched ellipse
by looking at tracking output; transforming the ellipse into a circle
using LU decomposition; then calculating the angle advanced.
"""
cwd = os.getcwd()
fout = open(self.output_filename + ".tmp", "w")
for i, closed_orbit in enumerate(self.closed_orbits_cached):
subs = closed_orbit["substitutions"]
for item, key in self.config.find_tune["subs_overrides"].items():
subs[item] = key
print("Finding tune with", end=' ')
for key in sorted(subs.keys()):
print(utilities.sub_to_name(key), subs[key], end=' ')
print()
tune_info = {"substitutions": subs}
ref_psv_list = self.track_decoupled(closed_orbit,
[0., 0., 0., 0.])[1:]
for axis, not_axis in [("u", "v"), ("v", "u")]:
finder = DPhiTuneFinder()
u, v = 0., 0.
if axis == "u":
u = self.config.find_tune["delta_1"]
else:
v = self.config.find_tune["delta_2"]
psv_list = self.track_decoupled(closed_orbit,
[u, 0., v, 0.])[1:]
self.set_finder(finder, axis, psv_list)
try:
tune = finder.get_tune(subs["__n_turns__"] /
10.) # tolerance
except:
tune = 0.
print(' Found', len(finder.dphi), 'dphi elements with tune',
tune, "+/-", finder.tune_error)
tune_info[axis + "_tune"] = tune
tune_info[axis + "_tune_rms"] = finder.tune_error
tune_info[axis + "_signal"] = list(zip(finder.u, finder.up))
tune_info[axis + "_dphi"] = finder.dphi
tune_info[axis + "_n_cells"] = len(finder.dphi)
canvas, hist, graph = finder.plot_phase_space_root()
name = os.path.join(
self.output_dir, "tune_" + str(i) + "_finding-" + axis +
"_" + axis + "_phase-space")
for format in ["png"]:
canvas.Print(name + "." + format)
canvas, hist, graph = finder.plot_cholesky_space_root()
name = os.path.join(
self.output_dir,
"tune_" + str(i) + "_" + axis + "_cholesky-space")
for format in ["png"]:
canvas.Print(name + "." + format)
self.set_finder(finder, not_axis, psv_list)
canvas, hist, graph = finder.plot_phase_space_root()
name = os.path.join(
self.output_dir, "tune_" + str(i) + "_finding-" + axis +
"_" + not_axis + "_phase-space")
for format in ["png"]:
canvas.Print(name + "." + format)
for i, u in enumerate([]): #finder.u[:-1]):
up = finder.up[i]
dphi = finder.dphi[i]
t = finder.t[i]
u_chol = finder.point_circles[i][0]
up_chol = finder.point_circles[i][1]
phi = math.atan2(up_chol, u_chol)
print(str(i).ljust(4), str(round(t, 4)).rjust(8), "...", \
str(round(u, 4)).rjust(8), str(round(up, 4)).rjust(8), "...", \
str(round(u_chol, 4)).rjust(8), str(round(up_chol, 4)).rjust(8), "...", \
str(round(phi, 4)).rjust(8), str(round(dphi, 4)).rjust(8))
for key in sorted(tune_info.keys()):
if "signal" not in key and "dphi" not in key:
print(" ", key, tune_info[key])
print(json.dumps(tune_info), file=fout)
fout.flush()
fout.close()
os.rename(self.output_filename + ".tmp", self.output_filename)
os.chdir(cwd)
def find_closed_orbit(sub_index, subs, seed, config):
"""
Find the closed orbit; algorithm is to track turn by turn; fit an ellipse to
the tracking; find the centre of the ellipse; repeat until no improvement or
10 iterations.
- energy: (float) kinetic energy at which the co is calculated
- step: (float) step size in tracking
- poly_order: (int) order of the polynomial fit to the field map (not used)
- smooth_oder: (int) order of smoothing the polynomial fit to the field map
(not used)
- seed: (list of 2 floats) [x, px] value to be used as the seed for the next
iteration; px value is ignored, sorry about that.
"""
max_iterations = config.find_closed_orbits["max_iterations"]
probe = config.find_closed_orbits["probe_files"]
for key in sorted(subs.keys()):
print(utilities.sub_to_name(key), subs[key], end=' ')
print()
out_dir = OUT_DIR
run_dir = RUN_DIR
tmp_dir = "./"
try:
os.makedirs(run_dir)
except OSError: # maybe the dir already exists
pass
os.chdir(run_dir)
print("Running in", os.getcwd())
common.substitute(CONFIG.tracking["lattice_file"],
tmp_dir + 'SectorFFAGMagnet.tmp', subs)
energy = subs["__energy__"]
ref_hit = reference(energy)
opal_exe = os.path.expandvars("${OPAL_EXE_PATH}/opal")
tracking = OpalTracking(tmp_dir + '/SectorFFAGMagnet.tmp',
tmp_dir + '/disttest.dat', ref_hit, probe,
opal_exe, tmp_dir + "/log")
seed_hit = ref_hit.deepcopy()
seed_hit["x"] = seed[0]
seed_hit["px"] = seed[1]
# fix momentum
seed_hit["pz"] = (ref_hit["p"]**2 - seed_hit["px"]**2)**0.5
print("Reference kinetic energy:", ref_hit["kinetic_energy"])
print("Seed kinetic energy: ", seed_hit["kinetic_energy"])
finder = EllipseClosedOrbitFinder(tracking, seed_hit)
generator = finder.find_closed_orbit_generator(["x", "px"], 1)
x_std_old = 1e9
i = -1
will_loop = True
iteration = None
while will_loop:
try:
iteration = next(generator)
except StopIteration:
will_loop = False
print(sys.exc_info()[1])
i += 1
#for i, iteration in enumerate(generator):
#for point in iteration.points:
heading = [
'station', 't', 'x', 'px', 'y', 'py', 'z', 'pz', 'r', 'pt',
'kinetic_energy'
]
for key in heading:
print(str(key).rjust(10), end=' ')
print()
for hit in tracking.last[0]:
for key in heading:
print(str(round(hit[key], 1)).rjust(10), end=' ')
print()
if iteration == None:
continue
print(iteration.centre)
#if iteration.centre != None: #i == 0 and
if i == 0:
plot_iteration(sub_index, i, iteration, energy)
if i >= max_iterations:
break
x_mean = numpy.mean([point[0] for point in iteration.points])
x_std = numpy.std([point[0] for point in iteration.points])
print("Seed:", iteration.points[0][0], "Mean:", x_mean, "Std:", x_std)
if type(iteration.centre) != type(
None) and x_std >= x_std_old: # require convergence
break
x_std_old = x_std
os.chdir(out_dir)
if i > 0:
plot_iteration(sub_index, i, iteration, energy)
return tracking.last[0]
