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 track_one(self, fields):
ref_probes = self.config.find_bump_parameters["ref_probe_files"]
#bump_probes = [self.config.find_bump_parameters["bump_probe_file"]]
energy = self.config.find_bump_parameters["energy"]
try:
os.makedirs(self.tmp_dir)
except OSError: # maybe the dir already exists
pass
self.setup_subs(fields)
os.chdir(self.tmp_dir)
find_closed_orbits.CONFIG = self.config
ref_hit = find_closed_orbits.reference(energy)
opal_exe = os.path.expandvars("${OPAL_EXE_PATH}/opal")
tracking = OpalTracking('SectorFFAGMagnet.tmp', 'disttest.dat',
ref_hit, ref_probes, opal_exe, "log")
test_hit = ref_hit.deepcopy()
closed_orbit = self.config.find_bump_parameters["closed_orbit"]
test_hit["x"] = closed_orbit[0]
test_hit["px"] = closed_orbit[1]
# fix momentum
test_hit["pz"] = (ref_hit["p"]**2 - test_hit["px"]**2)**0.5
print("Reference kinetic energy:", ref_hit["kinetic_energy"])
print("Seed kinetic energy: ", test_hit["kinetic_energy"])
hit_list = tracking.track_one(test_hit)
print("Station to probe mapping:\n ", end=' ')
for i, fname in enumerate(tracking.get_names()):
print("(" + str(i) + ",", fname + ")", end=' ')
print()
self.tracking_result = [[hit["station"], hit["x"], hit["px"]]
for hit in hit_list]
return hit_list
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 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_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 setup_tracking(config, probes, ref_energy):
ref_hit = reference(config, ref_energy)
opal_exe = os.path.expandvars(config.tracking["opal_path"])
lattice = config.tracking["lattice_file_out"]
log = config.tracking["tracking_log"]
beam = config.tracking["beam_file_out"]
tracking = OpalTracking(lattice, beam, ref_hit, probes, opal_exe, log)
return tracking
def setup_tracking(config, probes, ref_energy):
ref_hit = reference(config, ref_energy)
opal_exe = os.path.expandvars(config.tracking["opal_path"])
lattice = config.tracking["lattice_file_out"]
log = config.tracking["tracking_log"]
beam = config.tracking["beam_file_out"]
tracking = OpalTracking(lattice, beam, ref_hit, probes, opal_exe, log)
tracking.verbose = config.tracking["verbose"]
tracking.set_file_format(config.tracking["file_format"])
tracking.flags = config.tracking["flags"]
tracking.pass_through_analysis = StoreDataInMemory(config)
return tracking
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]
class DAFinder(object):
def __init__(self, config):
self.closed_orbit_file_name = os.path.join(
config.run_control["output_dir"],
config.find_closed_orbits["output_file"]) + ".out"
self.da_file_name = os.path.join(config.run_control["output_dir"],
config.find_da["get_output_file"])
self.scan_file_name = os.path.join(config.run_control["output_dir"],
config.find_da["scan_output_file"])
self.config = config
self.run_dir = os.path.join(config.run_control["output_dir"],
config.find_da["run_dir"])
self.co_list = self.load_closed_orbits()
self.ref_hit = None
self.min_delta = config.find_da["min_delta"]
self.max_delta = config.find_da["max_delta"]
self.data = []
self.fout_scan_tmp = None
self.fout_get_tmp = None
self.required_n_hits = config.find_da["required_n_hits"]
self.max_iterations = config.find_da["max_iterations"]
self.tracking = self.setup()
def get_all_da(self, co_index_list, seed_x, seed_y):
if co_index_list == None:
co_index_list = list(range(len(self.co_list)))
for i in co_index_list:
try:
co_element = self.co_list[i]
print("Finding da for element", i)
except KeyError:
print("Failed to find index", i, "in co_list of length",
len(co_list))
continue
if seed_x != None and seed_x > 0.:
co_element['x_da'] = self.get_da(co_element, 'x', seed_x)
if seed_y != None and seed_y > 0.:
co_element['y_da'] = self.get_da(co_element, 'y', seed_y)
print(json.dumps(co_element), file=self.fout_get())
self.fout_get().flush()
def da_all_scan(self, co_index_list, x_list, y_list):
if co_index_list == None:
co_index_list = list(range(len(self.co_list)))
for i in co_index_list:
try:
co_element = self.co_list[i]
print("Scanning da for element", i)
except KeyError:
print("Failed to find index", i, "in co_list of length",
len(co_list))
continue
co_element['da_scan'] = self.da_scan(co_element, x_list, y_list)
def load_closed_orbits(self):
fin = open(self.closed_orbit_file_name)
co_list = [json.loads(line) for line in fin.readlines()]
print("Loaded", len(co_list), "closed orbits")
return co_list
def setup(self):
self.tmp_dir = "./"
try:
os.makedirs(self.run_dir)
except OSError: # maybe the dir already exists
pass
os.chdir(self.run_dir)
print("Running in", os.getcwd())
self.opal_exe = os.path.expandvars("${OPAL_EXE_PATH}/opal")
def reference(self, hit_dict):
"""
Generate a reference particle
"""
hit = Hit.new_from_dict(hit_dict)
hit["x"] = 0.
hit["px"] = 0.
return hit
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 new_seed(self):
if self.data[-1][0] < self.min_delta: # reference run?
return None
if self.test_pass(
*self.data[-1]): # upper limit is okay; keep going up
if self.data[-1][0] > self.max_delta: # too big; give up
return None
return self.data[-1][0] * 2.
elif not self.test_pass(
*self.data[0]): # lower limit is bad; try going down
if abs(self.data[0][0]) < self.min_delta:
return None
return self.data[0][0] / 2.
else:
for i, item in enumerate(self.data[1:]):
if not self.test_pass(*item):
break
if abs(self.data[i][0] - self.data[i + 1][0]) < self.min_delta:
return None
return (self.data[i][0] + self.data[i + 1][0]) / 2.
def test_pass(self, seed, hits_list):
return len(hits_list) > self.required_n_hits
def events_generator(self, co_element, x_list, y_list):
co_hit = co_element["hits"][0]
for x in x_list:
for y in y_list:
a_hit = self.ref_hit.deepcopy()
a_hit['x'] = co_hit['x'] + x
a_hit['px'] = co_hit['px']
a_hit['y'] += y
yield {"x": x, "y": y}, a_hit
def da_scan(self, co_element, x_list, y_list):
self.setup_tracking(co_element)
gen = self.events_generator(co_element, x_list, y_list)
self.data = []
finished = False
while not finished:
event_list = []
track_list = []
try:
while len(event_list) < 1:
track, event = next(gen)
track_list.append(track)
event_list.append(event)
except StopIteration:
finished = True
if len(event_list) == 0:
break
many_tracks = self.tracking.track_many(event_list)
for i, hits in enumerate(many_tracks):
print("Tracked", len(hits), "total hits with track",
track_list[i], "first event x px y py", hits[0]["x"],
hits[0]["px"], hits[0]["y"], hits[0]["py"])
self.data.append(
[track_list[i], [a_hit.dict_from_hit() for a_hit in hits]])
print(json.dumps(self.data), file=self.fout_scan())
self.fout_scan().flush()
def get_da(self, co_element, axis, seed_x):
is_ref = abs(seed_x) < 1e-6
self.setup_tracking(co_element)
self.data = []
co_delta = {"x": 0, "y": 0}
iteration = 0
while seed_x != None and iteration < self.max_iterations:
co_delta[axis] = seed_x
my_time = time.time()
a_hit = Hit.new_from_dict(co_element["hits"][0])
a_hit[axis] += seed_x
try:
hits = self.tracking.track_one(a_hit)
except RuntimeError:
sys.excepthook(*sys.exc_info())
print("Never mind, keep on going...")
self.data.append(
[co_delta[axis], [a_hit.dict_from_hit() for a_hit in hits]])
self.data = sorted(self.data)
print("Axis", axis, "Seed", seed_x, "Number of cells hit",
len(hits), "in",
time.time() - my_time, "[s]")
sys.stdout.flush()
seed_x = self.new_seed()
if is_ref:
seed_x = None
iteration += 1
self.data = [list(item) for item in self.data]
return self.data
def fout_scan(self):
if self.fout_scan_tmp == None:
file_name = self.scan_file_name + ".tmp"
self.fout_scan_tmp = open(file_name, "w")
print("Opened file", file_name)
return self.fout_scan_tmp
def fout_get(self):
if self.fout_get_tmp == None:
file_name = self.da_file_name + ".tmp"
self.fout_get_tmp = open(file_name, "w")
print("Opened file", file_name)
return self.fout_get_tmp
class TrackBeam(object):
def __init__(self, config):
self.config = config
self.output_dir = config.run_control["output_dir"]
self.centre = None
self.ellipse = None
self.run_dir = ""
self.cwd = os.getcwd()
self.hits_in = []
self.hits_out = []
self.energy = None
def load_tune_data(self):
file_name = self.output_dir+"/"+self.config.find_tune["output_file"]
fin = open(file_name)
data = [json.loads(line) for line in fin.readlines()]
return data
def fit_tune_data(self, data):
eps_max = self.config.track_beam['eps_max']
x_emittance = self.config.track_beam['x_emittance']
y_emittance = self.config.track_beam['y_emittance']
sigma_pz = self.config.track_beam['sigma_pz']
sigma_z = self.config.track_beam['sigma_z']
self.energy = data['substitutions']['__energy__']
pid = self.config.tracking["pdg_pid"]
mass = xboa.common.pdg_pid_to_mass[abs(pid)]
p = ((self.energy+mass)**2 - mass**2)**0.5
x_centre, x_ellipse = xboa.common.fit_ellipse(
data['x_signal'],
eps_max,
verbose = False)
y_centre, y_ellipse = xboa.common.fit_ellipse(
data['y_signal'],
eps_max,
verbose = False)
x_ellipse *= (x_emittance/numpy.linalg.det(x_ellipse))**0.5
y_ellipse *= (y_emittance/numpy.linalg.det(y_ellipse))**0.5
self.centre = numpy.array([x for x in x_centre]+[y for y in y_centre]+[0., p])
self.ellipse = numpy.zeros((6, 6))
for i in range(2):
for j in range(2):
self.ellipse[i, j] = x_ellipse[i, j]
self.ellipse[i+2, j+2] = y_ellipse[i, j]
#for i in [0, 2]:
# self.ellipse[i+1, i+1] *= p*p
# self.ellipse[i, i+1] *= p
# self.ellipse[i+1, i] *= p
self.ellipse[4, 4] = sigma_z
self.ellipse[5, 5] = sigma_pz
print("Centre", self.centre)
print("Ellipse")
print(self.ellipse)
def setup_workspace(self):
self.run_dir = self.config.run_control["output_dir"]+"/"+self.config.track_beam["run_dir"]
try:
os.makedirs(self.run_dir)
except OSError:
pass # maybe the dir already exists
os.chdir(self.run_dir)
def generate_beam(self):
n_events = self.config.track_beam["subs_overrides"]["__n_events__"]
events = numpy.random.multivariate_normal(self.centre, self.ellipse, n_events)
keys = "x", "px", "y", "py", "z", "pz"
self.hits_in = []
for item in events:
hit = self.reference()
for i, key in enumerate(keys):
hit[key] = item[i]
self.hits_in.append(hit)
for hit in self.hits_in[0:10]:
print(" ", [hit[key] for key in keys])
print("Made", len(self.hits_in), "hits")
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)
#self.tracking._read_probes(phase_space_plots)
def reference(self):
"""
Generate a reference particle
"""
hit_dict = {}
hit_dict["pid"] = self.config.tracking["pdg_pid"]
hit_dict["mass"] = xboa.common.pdg_pid_to_mass[abs(hit_dict["pid"])]
hit_dict["charge"] = 1
hit_dict["x"] = 0.
hit_dict["kinetic_energy"] = self.energy
return Hit.new_from_dict(hit_dict, "pz")
def track(self):
try:
data = self.load_tune_data()
self.setup_workspace()
for i, item in enumerate(data):
self.fit_tune_data(item)
self.generate_beam()
self.run_tracking(i)
except:
raise
finally:
os.chdir(self.cwd)
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)
class DAFinder(object):
"""
DAFinder attempts to find the dynamic aperture by performing a binary search
of trajectories offset from the closed orbit. DA is determined to be the
trajectory that passes through at least a (user defined) number of probes.
"""
def __init__(self, config):
"""
Initialise the DAFinder object
- config: configuration object
"""
self.closed_orbit_file_name = os.path.join(config.run_control["output_dir"],
config.find_closed_orbits["output_file"])
self.da_file_name = os.path.join(config.run_control["output_dir"],
config.find_da["get_output_file"])
self.scan_file_name = os.path.join(config.run_control["output_dir"],
config.find_da["scan_output_file"])
self.config = config
self.run_dir = os.path.join(config.run_control["output_dir"],
config.find_da["run_dir"])
self.co_list = self.load_closed_orbits()
self.ref_hit = None
self.min_delta = config.find_da["min_delta"]
self.max_delta = config.find_da["max_delta"]
self.data = []
self.fout_scan_tmp = None
self.fout_get_tmp = None
self.required_n_hits = config.find_da["required_n_hits"]
self.max_iterations = config.find_da["max_iterations"]
self.tracking = self.setup()
def get_all_da(self, co_index_list, seed_x, seed_y):
"""
Get the DA
- co_index_list: list of indices to find DA for. Each element should be
an index from self.config.substitution_list. Set to
None to iterate over every element.
- seed_x: (float) best guess position offset for trajectory on the
horizontal DA. Set to None or negative value to disable
horizontal DA finding.
- seed_y: (float) best guess position offset for trajectory on the
vertical DA. Set to None or negative value to disable vertical
DA finding.
"""
if co_index_list == None:
co_index_list = list(range(len(self.co_list)))
for i in co_index_list:
try:
co_element = self.co_list[i]
print("Finding da for element", i)
except KeyError:
print("Failed to find index", i, "in co_list of length", len(co_list))
continue
if seed_x != None and seed_x > 0.:
co_element['x_da'] = self.get_da(co_element, 'x', seed_x)
if seed_y != None and seed_y > 0.:
co_element['y_da'] = self.get_da(co_element, 'y', seed_y)
print(json.dumps(co_element), file=self.fout_get())
self.fout_get().flush()
os.rename(self.da_file_name+".tmp", self.da_file_name)
def da_all_scan(self, co_index_list, x_list, y_list):
"""
Scan the DA
- co_index_list: list of indices to find DA for. Each element should be
an index from self.config.substitution_list. Set to
None to iterate over every element.
- seed_x: list of floats. Each list element is a horizontal position
offset from the closed orbit; the algorithm will track the
particle and count the number of probes through which the
particle passes.
- seed_y: list of floats. Each list element is a vertical position
offset from the closed orbit; the algorithm will track the
particle and count the number of probes through which the
particle passes.
The scan routine will generate a 2D grid in x and y. The trajectories
will be written to the scan file name.
"""
if co_index_list == None:
co_index_list = list(range(len(self.co_list)))
for i in co_index_list:
try:
co_element = self.co_list[i]
print("Scanning da for element", i)
except KeyError:
print("Failed to find index", i, "in co_list of length", len(co_list))
continue
co_element['da_scan'] = self.da_scan(co_element, x_list, y_list)
os.rename(self.scan_file_name+".tmp", self.scan_file_name)
def load_closed_orbits(self):
"""
Load the closed orbits file
"""
fin = open(self.closed_orbit_file_name)
co_list = [json.loads(line) for line in fin.readlines()]
print("Loaded", len(co_list), "closed orbits")
return co_list
def setup(self):
"""
Perform some setup
"""
self.tmp_dir = "./"
try:
os.makedirs(self.run_dir)
except OSError: # maybe the dir already exists
pass
os.chdir(self.run_dir)
print("Running in", os.getcwd())
self.opal_exe = os.path.expandvars("${OPAL_EXE_PATH}/opal")
def reference(self, hit_dict):
"""
Generate a reference particle
"""
hit = Hit.new_from_dict(hit_dict)
hit["x"] = 0.
hit["px"] = 0.
return hit
def setup_tracking(self, co_element):
"""
Setup the tracking routines
"""
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 new_seed(self):
"""
Generate a new seed for the DA finding routines.
* If all previous iterations have passed, then the largest offset is
doubled
* If all previous iterations have failed, then the smallest offset is
halved
* If some previous iterations have failed and some have passed, then a
binary interplation is performed between the highest passing
iteration and the lowest failing iteration
Returns the new seed or None if the iteration has finished.
"""
if self.data[-1][0] < self.min_delta: # reference run?
return None
if self.test_pass(*self.data[-1]): # upper limit is okay; keep going up
if self.data[-1][0] > self.max_delta: # too big; give up
return None
return self.data[-1][0]*2.
elif not self.test_pass(*self.data[0]): # lower limit is bad; try going down
if abs(self.data[0][0]) < self.min_delta:
return None
return self.data[0][0]/2.
else:
for i, item in enumerate(self.data[1:]):
if not self.test_pass(*item):
break
if abs(self.data[i][0]-self.data[i+1][0]) < self.min_delta:
return None
return (self.data[i][0]+self.data[i+1][0])/2.
def test_pass(self, seed, hits_list):
"""
Check to see if an iteration passed
Returns true if the number of hits in the hits list is greater than the
required number of hits.
"""
return len(hits_list) > self.required_n_hits
def events_generator(self, co_element, x_list, y_list):
"""
Generates the list of events for the da scan.
"""
co_hit = co_element["hits"][0]
for x in x_list:
for y in y_list:
a_hit = self.ref_hit.deepcopy()
a_hit['x'] = co_hit['x'] + x
a_hit['px'] = co_hit['px']
a_hit['y'] += y
yield {"x":x, "y":y}, a_hit
def da_scan(self, co_element, x_list, y_list):
"""
Do the da scan for a particular element of self.config.substitution_list.
"""
self.setup_tracking(co_element)
gen = self.events_generator(co_element, x_list, y_list)
self.data = []
finished = False
while not finished:
event_list = []
track_list = []
try:
while len(event_list) < 1:
track, event = next(gen)
track_list.append(track)
event_list.append(event)
except StopIteration:
finished = True
if len(event_list) == 0:
break
many_tracks = self.tracking.track_many(event_list)
for i, hits in enumerate(many_tracks):
print("Tracked", len(hits), "total hits with track", track_list[i], "first event x px y py", hits[0]["x"], hits[0]["px"], hits[0]["y"], hits[0]["py"])
self.data.append([track_list[i], [a_hit.dict_from_hit() for a_hit in hits]])
print(json.dumps(self.data), file=self.fout_scan())
self.fout_scan().flush()
def get_da(self, co_element, axis, seed_x):
"""
Do the da finding for a particular element of
self.config.substitution_list
"""
is_ref = abs(seed_x) < 1e-6
self.setup_tracking(co_element)
self.data = []
co_delta = {"x":0, "y":0}
iteration = 0
while seed_x != None and iteration < self.max_iterations:
co_delta[axis] = seed_x
my_time = time.time()
a_hit = Hit.new_from_dict(co_element["hits"][0])
a_hit[axis] += seed_x
try:
hits = self.tracking.track_one(a_hit)
except (RuntimeError, OSError):
sys.excepthook(*sys.exc_info())
print("Never mind, keep on going...")
hits = [a_hit]
self.data.append([co_delta[axis], [a_hit.dict_from_hit() for a_hit in hits]])
self.data = sorted(self.data)
print("Axis", axis, "Seed", seed_x, "Number of cells hit", len(hits), "in", time.time() - my_time, "[s]")
sys.stdout.flush()
seed_x = self.new_seed()
if is_ref:
seed_x = None
iteration += 1
self.data = [list(item) for item in self.data]
return self.data
def fout_scan(self):
"""
Open the scan output file
"""
if self.fout_scan_tmp == None:
file_name = self.scan_file_name+".tmp"
self.fout_scan_tmp = open(file_name, "w")
print("Opened file", file_name)
return self.fout_scan_tmp
def fout_get(self):
"""
Open the da finder output file
"""
if self.fout_get_tmp == None:
file_name = self.da_file_name+".tmp"
self.fout_get_tmp = open(file_name, "w")
print("Opened file", file_name)
return self.fout_get_tmp