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]