def find_closed_orbit(energy, nturns, step, poly_order, smooth_order, seed):
"""
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.
"""
print "Energy", energy, "NTurns", nturns, "StepSize", step, "Seed", seed, "Poly Order", poly_order, "Smooth Order", smooth_order
tmp_dir = "tmp/find_closed_orbits/"
try:
os.makedirs(tmp_dir)
except OSError: # maybe the dir already exists
pass
subs = {
'__energy__':energy,
'__stepsize__':step,
'__nturns__':nturns,
'__poly_order__':poly_order,
'__smooth_order__':smooth_order,
'__beamfile__':tmp_dir+'disttest.dat'
}
common.substitute('lattices/KurriMainRingTuneComparison/KurriMainRingTuneComparison.in', tmp_dir+'/Kurri_ADS_Ring.tmp', subs)
ref_hit = reference(energy)
opal_exe = os.path.expandvars("${OPAL_EXE_PATH}/opal")
tracking = OpalTracking(tmp_dir+'/Kurri_ADS_Ring.tmp', tmp_dir+'/disttest.dat', ref_hit, 'PROBE*.loss', opal_exe, tmp_dir+"/log")
mass = common.pdg_pid_to_mass[2212]
seed_hit = ref_hit.deepcopy()
seed_hit["x"] = seed[0]
finder = EllipseClosedOrbitFinder(tracking, seed_hit)
generator = finder.find_closed_orbit_generator(["x", "px"], 1)
x_std_old = 1e9
i = -1
for i, iteration in enumerate(generator):
print iteration.points
print iteration.centre
#if iteration.centre != None: #i == 0 and
if i == 0:
plot_iteration(i, iteration, energy, step, poly_order, smooth_order)
if i >= 10:
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 iteration.centre != None and x_std >= x_std_old: # require convergence
break
x_std_old = x_std
if i > -1:
plot_iteration(i, iteration, energy, step, poly_order, smooth_order)
return tracking.last[0]
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 tracking(self, n_particles):
#hit_list = []
#for i, amplitude in enumerate([1.]+[i*10. for i in range(1, 11)]):
# hit_list += self.make_bunch(amplitude)
common.substitute("config_gaussian.in", "tmp/config.py",
{
"__n_particles__":n_particles,
})
#bunch = Bunch.new_from_hits(hit_list)
#bunch.hit_write_builtin("icool_for003", "tmp/for003.dat")
sim = os.path.expandvars("${MAUS_ROOT_DIR}/bin/simulate_mice.py")
proc = subprocess.Popen(['python', sim, '--configuration_file', 'tmp/config.py'])
proc.wait()
def tracking(self, n_particles):
#hit_list = []
#for i, amplitude in enumerate([1.]+[i*10. for i in range(1, 11)]):
# hit_list += self.make_bunch(amplitude)
common.substitute("config_gaussian.in", "tmp/config.py", {
"__n_particles__": n_particles,
})
#bunch = Bunch.new_from_hits(hit_list)
#bunch.hit_write_builtin("icool_for003", "tmp/for003.dat")
sim = os.path.expandvars("${MAUS_ROOT_DIR}/bin/simulate_mice.py")
proc = subprocess.Popen(
['python', sim, '--configuration_file', 'tmp/config.py'])
proc.wait()
def common_substitute_test():
out1 = open('test1.in', 'w')
if not out1: return 'warn'
out1.write('some test input\nis written here')
out1.close()
common.substitute('test1.in', 'test1.out', {'some':'a bit of', 'test':'tested', 'input':'output', 'written':'read', 'there':'here', 'kx':'uga'})
out2 = open('test1.out', 'r')
if not out2: return 'warn'
lines = out2.readlines()
out2.close()
if not lines == ['a bit of tested output\n','is read here']: return 'fail'
os.remove('test1.in')
os.remove('test1.out')
return 'pass'
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 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]
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 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_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.
"""
fout = open(self.output, "w")
index = 0
for energy, position in sorted(self.closed_orbits_cached.iteritems()):
if self.energy_min != None and self.energy_max != None:
if energy < self.energy_min or energy > self.energy_max:
continue
index += 1
if index > 2:
ROOT.gROOT.SetBatch(True)
print "Finding tune at", energy, "MeV and closed orbit x=", position, "mm"
self.co_x = position
tune_info = {
"energy":energy,
"nturns":self.nturns,
"stepsize":self.step_size,
"poly_order":self.poly_order,
"smooth_order":self.smooth_order,
}
for axis1, axis2, delta1, delta2 in [("x", "px", self.delta_x, 0.),
("y", "py", self.delta_y, 0.)]:
hit = self._reference(energy)
hit["x"] = position
self._temp_dir()
common.substitute(
self.lattice_src,
self.tmp_dir+self.lattice, {
"__energy__":energy,
"__nturns__":self.nturns,
"__beamfile__":self.tmp_dir+self.beam_file,
"__stepsize__":self.step_size,
"__poly_order__":self.poly_order,
"__smooth_order__":self.smooth_order,
})
tracking = OpalTracking(self.tmp_dir+self.lattice,
self.tmp_dir+self.beam_file,
self._reference(energy),
self.output_filename,
self.opal,
self.tmp_dir+self.log_file)
finder = DPhiTuneFinder()
finder.run_tracking(axis1, axis2, delta1, delta2, hit, tracking)
for track_index, track in enumerate(tracking.last):
print 'Track', track_index, 'of', len(tracking.last), \
'with', len(track), 'hits'
for hit in track:
print ' ', hit['t'], 'polar:', math.atan2(hit['y'], \
hit['x']), (hit['y']**2+hit['x']**2)**0.5, \
'cart:', hit['x'], hit['y'], hit['z']
tune = finder.get_tune(self.nturns/10.)
tune_info[axis1+"_tune"] = tune
tune_info[axis1+"_tune_error"] = finder.tune_error
tune_info[axis1+"_signal"] = zip(finder.u, finder.up)
for key in sorted(tune_info.keys()):
if "signal" not in key:
print " ", key, tune_info[key]
print >> fout, json.dumps(tune_info)
fout.flush()
