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 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 __init__(self, config):
"""
Find the tune.
- probe_file_name: name of a PROBE file from OPAL output, or None. If
file_name is specified, use that file_name in order to
calculate tunes (and generate plots), otherwise generate a
new one by tracking.
- closed_orbits_file_name: name of a file containing closed orbits,
generated by e.g.
"""
self.closed_orbits_cached = None # filled by _load_closed_orbits
self.tmp_dir = None
self.unique_id = 1
self.lattice_src = config.tracking["lattice_file"]
self.probe_filename = config.find_tune["probe_files"]
co_file = os.path.join(config.run_control["output_dir"],
config.find_closed_orbits["output_file"])
self._load_closed_orbits(co_file)
self.delta_x = config.find_tune["delta_x"]
self.delta_y = config.find_tune["delta_y"]
self.row = config.find_tune["row_list"]
self.do_axis = config.find_tune["axis"]
self.opal = config.tracking["opal_path"]
self.step_size = config.tracking["step_size"]
self.config = config
self.lattice = "/SectorFFAGMagnet.tmp"
self.beam_file = "/disttest.dat"
self.log_file = "/log"
self.output_dir = config.run_control["output_dir"]
self.output_filename = os.path.join(self.output_dir,
config.find_tune["output_file"])
self.plot_dir = os.path.join(self.output_dir, "find_tune")
utilities.clear_dir(self.plot_dir)
def main(file_list):
flip_vertical = False
number_of_phases = 1
score_tolerance = 1000.
foil_probe = 0
foil_var = 3
foil_axis = "Height at foil [mm]"
data = []
plot_dir = os.path.split(file_list[0])[0] + "/plot_bump"
utilities.clear_dir(plot_dir)
index = 0
score, n_iterations = [], []
for file_glob in file_list:
for file_name in sorted(glob.glob(file_glob)):
index += 1
if index % number_of_phases != 0:
score.append(
load_file_alt(file_name, flip_vertical)[0]["score"])
n_iterations.append(
load_file_alt(file_name, flip_vertical)[0]["n_iterations"])
continue
data += load_file_alt(file_name, flip_vertical)
data[-1]["score"] = score + [data[-1]["score"]]
data[-1]["n_iterations"] = n_iterations + [
data[-1]["n_iterations"]
]
score, n_iterations = [], []
traj_fname = "output/triplet_baseline/anticorrelated_painting/toy_model/dp_p_0.0013__col_dens_2e-05__n_tot_50__n_i_25__inj_emit_0.1__k_1.0__tgt_emit_8.0/run_summary.json"
trajectory = [
] #load_trajectory(traj_fname, 0, [3744.07267-10, 3.843, -89.83+8, -1.158], 32, 75.0)
for probe in [0, 7]:
for axis in ['x', 'y', 'x-y']:
if probe == 0:
add_points = trajectory + \
[[probe]+ [3744.07267, 3.843, -89.83, -1.158],] #3744.08063, 3.843, -9.81, 1.151],)
else:
add_points = ([probe] + [
3739.427804804206, -0.0002874136417290174,
-88.77890374233482, -0.0008126002995965109
], )
if flip_vertical:
for point in add_points:
point[2] *= -1
point[3] *= -1
plot_closed_orbit(data, axis, plot_dir, probe, add_points)
plot_fields(data, plot_dir, foil_probe, foil_var, foil_axis)
plot_fields_2(data, plot_dir, [0, 1, 2, 3, 4])
plot_all_fields_2d(data, plot_dir, foil_probe, 3, 1)
for bump in data:
print(bump["score"])
plot_phased_var_2d(data, plot_dir, foil_probe, 3, 1, "score",
number_of_phases)
plot_phased_var_2d(data, plot_dir, foil_probe, 3, 1, "n_iterations",
number_of_phases)
def parse_args(self, args):
parser = argparse.ArgumentParser()
parser.add_argument('file_names', metavar='N', type=str, nargs='+')
parser.add_argument('--log_x', dest='log_x', action='store_true')
parser.set_defaults(log_x=False)
args = parser.parse_args()
self.output_dir = os.path.split(args.file_names[0])[0]
self.output_dir = os.path.join(self.output_dir, "plot_closed_orbits")
utilities.clear_dir(self.output_dir)
for file_name in args.file_names:
self.load_file(file_name)
self.log_x = args.log_x
def main(config):
global CONFIG, OUT_DIR, RUN_DIR
CONFIG = config
OUT_DIR = CONFIG.run_control["output_dir"]
utilities.clear_dir(os.path.join(OUT_DIR, "find_closed_orbits"))
RUN_DIR = os.path.join(OUT_DIR, "tmp/find_closed_orbits/")
fout_name = os.path.join(OUT_DIR, config.find_closed_orbits["output_file"])
fout = open(fout_name + ".tmp", 'w')
subs_list = config.substitution_list
results = []
for i, sub in enumerate(subs_list):
will_loop = True
is_batch = i >= config.find_closed_orbits["root_batch"]
for item, key in config.find_closed_orbits["subs_overrides"].items():
sub[item] = key
ROOT.gROOT.SetBatch(is_batch)
hit_list = []
while will_loop:
try:
next_seed, will_loop = get_seed(config, results, sub)
except Exception:
co_hit = results[0]["hits"][0]
next_seed = [co_hit["x"], co_hit["px"]]
will_loop = False
print("LOOPING DEBUG", next_seed, will_loop, sub['__energy__'])
try:
hit_list = find_closed_orbit(i, sub, next_seed, config)
except IndexError as ValueError:
sys.excepthook(*sys.exc_info())
except RuntimeError:
sys.excepthook(*sys.exc_info())
hit_list = []
print("Breaking loop due to tracking error")
will_loop = False
will_loop = len(hit_list) < 20 and will_loop
output = {
"substitutions": sub,
"hits": [hit.dict_from_hit() for hit in hit_list]
}
results.append(output)
print(json.dumps(output), file=fout)
fout.flush()
if config.find_closed_orbits["do_plot_orbit"] and i == 0:
plot_orbit.main(config.run_control["output_dir"],
"tmp/find_closed_orbits/",
"SectorFFAGMagnet-trackOrbit.dat")
fout.close()
time.sleep(1)
os.rename(fout_name + ".tmp", fout_name + ".out")
ROOT.gROOT.SetBatch(False)
def find_all_closed_orbits(self):
"""
Outer loop of the closed orbit finder
For each substitution in the config.substitution_list, attempt to find
the closed orbit. Write the closed orbits into the output json file.
"""
utilities.clear_dir(os.path.join(self.out_dir, "find_closed_orbits"))
fout_name = os.path.join(self.out_dir,
self.config.find_closed_orbits["output_file"])
fout = open(fout_name + ".tmp", 'w')
subs_list = self.config.substitution_list
results = []
for i, sub in enumerate(subs_list):
will_loop = True
is_batch = i >= self.config.find_closed_orbits["root_batch"]
for item, key in self.config.find_closed_orbits[
"subs_overrides"].items():
sub[item] = key
hit_list = []
while will_loop:
try:
next_seed, will_loop = self.get_seed(results, sub)
except Exception:
co_hit = results[0]["hits"][0]
next_seed = [co_hit["x"], co_hit["px"]]
will_loop = False
try:
hit_list = self.find_closed_orbit(i, sub, next_seed)
except IndexError as ValueError:
sys.excepthook(*sys.exc_info())
except RuntimeError:
sys.excepthook(*sys.exc_info())
hit_list = []
print("Breaking loop due to tracking error")
will_loop = False
will_loop = len(hit_list) < 20 and will_loop
output = {
"substitutions": sub,
"hits": [hit.dict_from_hit() for hit in hit_list]
}
results.append(output)
print(json.dumps(output), file=fout)
fout.flush()
fout.close()
time.sleep(1)
os.rename(fout_name + ".tmp", fout_name)
def plot_da(data, max_n_points, plot_dir, acceptance):
variables = utilities.get_substitutions_axis(data)
plot_dir = os.path.join(plot_dir, "plot_da2")
utilities.clear_dir(plot_dir)
for index, item in enumerate(data):
for key in variables:
variables[key] = item['substitutions'][key]
for da_key in 'y_da', 'x_da':
if da_key not in list(item.keys()):
print("Did not find", da_key, "in keys", list(item.keys()),
"... skipping")
continue
for axis1, axis2 in ('x', "x'"), ('y', "y'"):
canvas = plot_one_da(item, da_key, axis1, axis2, max_n_points,
variables, acceptance)
plot_name = "da_" + str(
index) + "_" + da_key + "_" + axis1 + "_" + axis2
for format in ["eps", "png", "root"]:
canvas.Print(plot_dir + "/" + plot_name + "." + format)
if axis1[0] not in da_key:
continue
def plot_da(self, max_n_points, plot_dir, acceptance):
variables = utilities.get_substitutions_axis(self.data,
"substitutions")
plot_dir = os.path.join(plot_dir, "plot_da")
utilities.clear_dir(plot_dir)
self.plot_dir = plot_dir
for index, item in enumerate(self.data):
for key in variables:
variables[key] = item['substitutions'][key]
for da_key in 'y_da', 'x_da':
if da_key not in list(item.keys()):
print("Did not find", da_key, "in keys", list(item.keys()),
"... skipping")
continue
for axis1, axis2 in ('x', "px"), ('y', "py"):
canvas = self.plot_one_da(item, da_key, axis1, axis2,
max_n_points, variables,
acceptance)
plot_name = "da_" + str(
index) + "_" + da_key + "_" + axis1 + "_" + axis2
for format in ["eps", "png", "root"]:
canvas.Print(plot_dir + "/" + plot_name + "." + format)
continue
if axis1[0] not in da_key:
continue
canvas, chol_canvas = plot_one_phi(item, da_key, axis1,
axis2, max_n_points,
variables)
plot_name = "phi_" + str(
index) + "_" + da_key + "_" + axis1 + "_" + axis2
for format in ["eps", "png", "root"]:
canvas.Print(plot_dir + "/" + plot_name + "." + format)
chol_canvas.Print(plot_dir + "/" + plot_name +
"_cholesky." + format)
self.plot_decoupled_da(item, da_key, index, axis1, axis2,
max_n_points, variables, acceptance)
def __init__(self, config):
"""
Find the tune.
-probe_file_name: name of a PROBE file from OPAL output, or None. If file_name
is specified, use that file_name in order to calculate tunes
(and generate plots), otherwise generate a new one by
tracking.
-closed_orbits_file_name: name of a file containing closed orbits,
generated by e.g.
"""
self.config = config
self.var_list = ["x", "px", "y", "py"]
self.output_dir = os.path.join(self.config.run_control["output_dir"])
self.output_filename = os.path.join(
self.output_dir, self.config.find_tune["output_file"])
self.run_dir = os.path.join(self.output_dir,
self.config.find_tune["run_dir"])
self._load_closed_orbits()
self.tracking = None
self.here = os.getcwd()
DecoupledTransferMatrix.det_tolerance = 1.
utilities.clear_dir(self.run_dir)
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 setup_output_dir(self):
if self.output_dir == "":
self.output_dir = os.path.split(
self.closed_orbit)[0] + "/toy_model/"
utilities.clear_dir(self.output_dir)