def track_painting(self, bump_list):
foil_probe = self.config.track_bump["foil_probe_files"]
ramp_probe = self.config.track_bump["ramp_probe_files"]
foil_phi = self.config.track_bump["foil_probe_phi"]
ramp_phi = self.config.track_bump["ramp_probe_phi"]
injected_beam = self.injection_orbit
energy = self.config.track_bump["energy"]
self.do_substitutions({}, 0, +1.0)
foil_tracking = utilities.setup_tracking(self.config, foil_probe, energy)
ramp_tracking = utilities.setup_tracking(self.config, ramp_probe, energy)
foil_tracking.verbose = False
ramp_tracking.verbose = False
test_hit = utilities.reference(self.config, energy, injected_beam[0],
injected_beam[1], injected_beam[2],
injected_beam[3])
print("Seed kinetic energy: ", test_hit["kinetic_energy"])
# fix momentum
hit_list_in = []
log = open("track_bump.log", "w")
for bump_fields, n_turns, closed_orbit in bump_list:
# ramp the fields
print("Ramping fields", bump_fields, file=log)
if len(hit_list_in) > 0:
# track from the ramp point to the foil point
print("Tracking to foil", file=log)
self.do_substitutions(bump_fields, ramp_phi, +1.0)
hit_list_of_lists = foil_tracking.track_many(hit_list_in)
hit_list_in = [hit_list[-1] for hit_list in hit_list_of_lists]
self.do_substitutions(bump_fields, foil_phi, +1.0)
for i in range(n_turns):
hit_list_in.append(copy.deepcopy(test_hit))
hit_list_in = [copy.deepcopy(test_hit),
copy.deepcopy(test_hit)] + hit_list_in
hit_list_of_lists = foil_tracking.track_many(hit_list_in)
hit_list_of_lists = hit_list_of_lists[2:]
hit_list_in = []
for hit_list in hit_list_of_lists:
if hit_list[-1]["t"] > 100.:
hit_list_in.append(hit_list[-1])
self.output.append(hit_list_in)
print(" After turn", i + 1, "fields", end=' ', file=log)
for field in bump_fields:
print(field, end=' ')
print("found", len(hit_list_of_lists), "particles", file=log)
for item in hit_list_in:
print(" ", \
format(item["station"], "6"), \
format(item["t"], "12.6g"), \
format(item["x"], "12.6g"), \
format(item["px"], "12.6g"), \
format(item["kinetic_energy"], "12.6g"), file=log)
# now track from the foil to the ramp point
print("Tracking to ramp point", file=log)
hit_list_of_lists = ramp_tracking.track_many(hit_list_in)
hit_list_in = [hit_list[-1] for hit_list in hit_list_of_lists]
return hit_list
def tracking(self):
probe_files = self.config.find_fixed_frequency["probe_files"]
phi_list = [i / 10.
for i in range(-5, 6)] + [-0.05, -0.01, +0.05, +0.01]
phi_list = sorted(phi_list, key=lambda x: abs(x))
clr_list = numpy.array([
1, ROOT.kRed + 2, ROOT.kRed, ROOT.kRed - 9, ROOT.kBlue + 2,
ROOT.kBlue, ROOT.kBlue - 9
])
ldots = [{'kinetic_energy': '...', 't': '...'}]
for i, phi in enumerate(phi_list):
self.do_substitutions(phi) # RF phase goes from 0 to 1
test_hit = self.get_hit()
tracking = utilities.setup_tracking(self.config, probe_files,
test_hit['kinetic_energy'])
os.chdir(self.tmp_dir)
hit_list = tracking.track_one(test_hit)
lattice_name = self.config.tracking["lattice_file_out"].split(
".")[0]
os.rename(lattice_name + "-trackOrbit.dat",
lattice_name + "-trackOrbit-" + str(i + 1) + ".dat")
max_e = max([hit['kinetic_energy'] for hit in hit_list])
print(len(hit_list), "hits with peak energy", max_e,
"start and end t, E:")
print([(hit['t'], hit['kinetic_energy'])
for hit in hit_list[0:2] + ldots + hit_list[-3:-1]])
color = clr_list[0]
clr_list = numpy.roll(clr_list, 1)
self.plot_acceleration(hit_list, phi + 0.5, color)
def track_one(self, fields):
utilities.clear_dir(self.tmp_dir)
os.chdir(self.tmp_dir)
ref_probes = self.config.find_bump_parameters["ref_probe_files"]
#bump_probes = [self.config.find_bump_parameters["bump_probe_file"]]
ref_energy = self.config.find_bump_parameters["energy"]
self.setup_subs(fields)
tracking = utilities.setup_tracking(self.config, ref_probes,
ref_energy)
test_hit = tracking.ref.deepcopy()
closed_orbit = self.config.find_bump_parameters["closed_orbit"]
test_hit["x"] = closed_orbit[0]
test_hit["px"] = closed_orbit[1]
test_hit["y"] = closed_orbit[2]
test_hit["py"] = closed_orbit[3]
# fix momentum
test_hit["pz"] = (tracking.ref["p"]**2 - test_hit["px"]**2)**0.5
print("Reference kinetic energy:", tracking.ref["kinetic_energy"])
print("Seed kinetic energy: ",
test_hit["kinetic_energy"],
flush=True)
hit_list = tracking.track_many([test_hit])[1]
print("Station to probe mapping:\n ", end=' ')
for i, fname in enumerate(tracking.get_name_list()):
print("(" + str(i) + ",", fname + ")", end=' ')
print()
hit_list = self.cuts(hit_list)
self.tracking_result = [[
hit["station"], hit["x"], hit["px"], hit["y"], hit["py"]
] for hit in hit_list]
return hit_list
def track_one(self, subs_overrides):
utilities.clear_dir(self.tmp_dir)
os.chdir(self.tmp_dir)
probes = self.config.build_bump_surrogate_model["ref_probe_files"]
ref_energy = self.config.build_bump_surrogate_model["energy"]
tracking = utilities.setup_tracking(self.config, probes, ref_energy)
test_hit = tracking.ref.deepcopy()
closed_orbit = self.config.build_bump_surrogate_model["closed_orbit"]
test_hit["x"] = closed_orbit[0]
test_hit["px"] = closed_orbit[1]
test_hit["y"] = closed_orbit[2]
test_hit["py"] = closed_orbit[3]
# fix momentum
test_hit["pz"] = (tracking.ref["p"]**2-test_hit["px"]**2)**0.5
#print("Reference kinetic energy:", tracking.ref["kinetic_energy"])
#print("Seed kinetic energy: ", test_hit["kinetic_energy"], flush=True)
if not self.optimisation["fore_tracking"]:
test_hit["px"] *= -1
test_hit["py"] *= -1
test_hit["pz"] *= -1
subs_overrides["__beam_charge__"] = -1
utilities.do_lattice(self.config, self.subs, subs_overrides)
if self.dummy_run: # dummy run
fields = [2*subs_overrides[fname] for fname in self.get_fields()]
self.tracking_result = [[i]+fields for i in range(10)]
else:
hit_list = tracking.track_many([test_hit])[1]
self.tracking_result = [[hit["station"], hit["x"], hit["px"], hit["y"], hit["py"]] for hit in hit_list]
def setup_tracking(self, co_element):
overrides = self.config.find_da["subs_overrides"]
overrides["__n_particles__"] = 3
utilities.do_lattice(self.config, co_element["substitutions"],
overrides)
self.tracking = utilities.setup_tracking(
self.config, self.config.find_da["probe_files"],
co_element["substitutions"]["__energy__"])
def fore_track(self, fields, proton_orbit):
self.file_index += 1
foil_phi = self.config.track_bump["proton_orbit_phi"]
foil_probe = self.config.track_bump["foil_probe_files"]
energy = self.config.track_bump["energy"]
self.do_substitutions(fields, foil_phi, +1.0)
tracking = utilities.setup_tracking(self.config, foil_probe, energy)
test_hit = utilities.reference(self.config, energy, proton_orbit[0],
proton_orbit[1], proton_orbit[2],
proton_orbit[3])
try:
tracking.track_many([test_hit])
except OSError: # did not load any output files - we don't care
pass
lattice_name = self.config.tracking["lattice_file_out"].split(".")[0]
os.rename(
lattice_name + "-trackOrbit.dat",
lattice_name + "-trackOrbit-fore-" + str(self.file_index) + ".dat")
def track_decoupled(self, co_element, decoupled_psv):
overrides = self.config.find_tune["subs_overrides"]
tm = copy.deepcopy(co_element["tm"])
for i, row in enumerate(tm):
tm[i] = row[1:5]
tm = DecoupledTransferMatrix(tm)
coupled_psv = tm.coupled(decoupled_psv)
print("Decoupled:", decoupled_psv)
print("Coupled: ", coupled_psv)
ref_hit = Hit.new_from_dict(co_element["seed_hit"])
test_hit = copy.deepcopy(ref_hit)
for i, var in enumerate(self.var_list):
print("Adding",
coupled_psv[i],
"to",
var,
"before",
test_hit[var],
end=' ')
test_hit[var] += coupled_psv[i]
print("after", test_hit[var])
test_hit.mass_shell_condition("pz")
os.chdir(self.run_dir)
subs = utilities.do_lattice(self.config, co_element["substitutions"],
overrides)
self.tracking = utilities.setup_tracking(
self.config, self.config.find_tune["probe_files"],
subs["__energy__"])
self.tracking.verbose = True
track = self.tracking.track_one(test_hit)
print("Found", len(track), "hits")
os.chdir(self.here)
decoupled_track = []
for hit in track:
print(" ", end=' ')
for var in self.var_list:
print(format(hit[var], "14.8g"), end=' ')
print()
coupled_psv_out = [
hit[var] - ref_hit[var] for var in self.var_list
]
decoupled_psv_out = tm.decoupled(coupled_psv_out)
decoupled_track.append(decoupled_psv_out)
return decoupled_track
def back_track(self, fields):
self.file_index += 1
foil_phi = self.config.track_bump["foil_phi"]
injected_beam = self.injection_orbit
foil_probe = self.config.track_bump["foil_probe_files"]
energy = self.config.track_bump["energy"]
self.do_substitutions(fields, foil_phi, +1.0)
tracking = utilities.setup_tracking(self.config, foil_probe, energy)
test_hit = utilities.reference(self.config, energy, injected_beam[0],
injected_beam[1], injected_beam[2],
injected_beam[3])
test_hit["px"] *= -1
test_hit["py"] *= -1
test_hit["pz"] *= -1
try:
tracking.track_many([test_hit])
except OSError: # did not load any output files - we don't care
pass
lattice_name = self.config.tracking["lattice_file_out"].split(".")[0]
os.rename(
lattice_name + "-trackOrbit.dat",
lattice_name + "-trackOrbit-back-" + str(self.file_index) + ".dat")
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]