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 test_check_co_json_repr(self):
"""
test EllipseClosedOrbitFinderIteration.json_repr
"""
n_turns = 10
co_finder = EllipseClosedOrbitFinder(self.tracking, self.co)
test_values = {'x': 11., 'px': 8., 'y': 1., 'py': 1.}
co_check = co_finder.check_closed_orbit(test_values, n_turns)
self.assertEqual(co_check.json_repr()["points"], co_check.points)
self.assertEqual(co_check.json_repr()["keys"], co_check.keys)
self.assertEqual(co_check.json_repr()["eps_max"], co_check.eps_max)
def test_check_co_t(self):
"""test EllipseClosedOrbitFinder.check_co - t"""
n_turns = 5
co_finder = EllipseClosedOrbitFinder(self.tracking, self.co)
test_values = {'t': 0.1, 'energy': 1000.1}
co_check = co_finder.check_closed_orbit(test_values, n_turns)
# Should be divergent
# co_check.plot_ellipse('t', 'energy', 'ns', 'MeV')
# note this makes a more-or-less straight line in t-energy so tends to
# produce singular ellipses
self.assertGreater(co_check.get_sigma_noise() * 2.,
co_check.get_mean_noise())
def test_check_co_xy(self):
"""test EllipseClosedOrbitFinder.check_co - xy
"""
n_turns = 10
co_finder = EllipseClosedOrbitFinder(self.tracking, self.co)
test_values = {'x': 11., 'px': 8., 'y': 1., 'py': 1.}
co_check = co_finder.check_closed_orbit(test_values, n_turns)
# check that we have converged in 2D
for i, key in enumerate(co_check.keys):
self.assertLess(abs(self.co[key] - co_check.centre[i]),
abs(self.co[key] - test_values[key]) / 2.)
self.assertLess(co_check.get_sigma_noise() * 2.,
co_check.get_mean_noise())
def test_check_co_y(self):
"""test EllipseClosedOrbitFinder.check_co - y"""
n_turns = 10
co_finder = EllipseClosedOrbitFinder(self.tracking, self.co)
test_values = {'y': 1., 'py': 1.}
co_check = co_finder.check_closed_orbit(test_values, n_turns)
# check we have reasonable convergence properties
for i, key in enumerate(co_check.keys):
self.assertLess(abs(self.co[key] - co_check.centre[i]),
abs(self.co[key] - test_values[key]) / 2.)
# check noise is of correct length
n_noise = float(len(co_check.noise))
self.assertGreater(n_noise, n_turns - 1)
# noise should be "clean" if we are fitting the ellipse properly
self.assertLess(co_check.get_sigma_noise() * 2.,
co_check.get_mean_noise())
def test_check_co_plot_ellipse(self):
"""
test EllipseClosedOrbitFinderIteration.json_repr
"""
n_turns = 20
co_finder = EllipseClosedOrbitFinder(self.tracking, self.co)
test_values = {'x': 11., 'px': 8., 'y': 1., 'py': 1.}
co_check = co_finder.check_closed_orbit(test_values, n_turns)
canvas, hist, ellipse, graph = \
co_check.plot_ellipse("x", "px", "mm", "MeV/c", 4, "test_fit")
# mock-up that we failed to find an ellipse; check that doesn't throw
co_check.centre = None
co_check.ellipse = None
co_check.noise = None
canvas, hist, ellipse, graph = \
co_check.plot_ellipse("x", "px", "mm", "MeV/c", 4, "no_ell")
def test_scan_data(self):
"""
check that the scan data is calculated correctly
"""
co_finder = EllipseClosedOrbitFinder(self.tracking, self.co)
scan_start = {'x': 9., 'px': 6.}
scan_end = {'x': 11.1, 'px': 8.1}
scan_step = {'x': 1., 'px': 1.}
scan_gen = co_finder.scan_generator(scan_start, scan_end, scan_step)
co_iter = scan_gen.next()
self.assertEqual(len(co_iter.points), 6)
co_iter = scan_gen.next()
co_iter = scan_gen.next()
co_iter = scan_gen.next()
# this should be on the closed orbit
co_iter = scan_gen.next()
for point in co_iter.points:
for i, x in enumerate(point):
self.assertAlmostEqual(point[i], [7., 10.][i])
def test_check_co_x(self):
"""test EllipseClosedOrbitFinder.check_co - x"""
n_turns = 10
co_finder = EllipseClosedOrbitFinder(self.tracking, self.co)
test_values = {'x': 11., 'px': 8.}
co_check = co_finder.check_closed_orbit(test_values, n_turns)
self.assertEqual(co_check.keys, sorted(test_values.keys()))
# check we have reasonable convergence properties
for i, key in enumerate(co_check.keys):
self.assertLess(abs(self.co[key] - co_check.centre[i]),
abs(self.co[key] - test_values[key]) / 2.)
# check noise is of correct length
n_noise = float(len(co_check.noise))
self.assertGreater(n_noise, n_turns - 1)
# noise should be "clean" if we are fitting the ellipse properly
mean_noise = sum(co_check.noise) / n_noise
self.assertAlmostEqual(co_check.get_mean_noise(), mean_noise)
noise_squ = [noise * noise for noise in co_check.noise]
sigma_noise = (sum(noise_squ) / n_noise - mean_noise**2)**0.5
self.assertAlmostEqual(co_check.get_sigma_noise(), sigma_noise)
def test_scan_steps(self):
"""
test EllipseClosedOrbitFinder.scan_generator generates correct start x
"""
co_finder = EllipseClosedOrbitFinder(self.tracking, self.co)
scan_start = {'x': -1., 'y': -2., 'z': -3.}
scan_end = {'x': 4.1, 'y': 5., 'z': 6.433}
scan_step = {'x': 1., 'y': 2., 'z': 3.}
scan_gen = co_finder.scan_generator(scan_start, scan_end, scan_step)
start = [scan_start[key] for key in ['x', 'y', 'z']]
end = [scan_end[key] for key in ['x', 'y', 'z']]
last_step = scan_gen.next().points[0]
self.assertEqual(last_step, start)
for item in scan_gen:
step = item.points[0]
self.assertGreater(step, last_step)
self.assertLess(step, end)
last_step = step
scan_gen = co_finder.scan_generator(scan_start, scan_end, scan_step)
scan_list = [x for x in scan_gen]
self.assertEqual(len(scan_list), 6 * 4 * 4)
def test_co_finder(self):
"""
Check that the co finder converges
"""
test_hit = self.co.deepcopy()
test_hit["x"] += 1.
test_hit["y"] += 1.
co_finder = EllipseClosedOrbitFinder(self.tracking, test_hit)
for test_keys in [["x", "px"], ["x", "px", "y", "py"]]:
co_generator = co_finder.find_closed_orbit_generator(test_keys, 10)
try:
canv = None
for my_co in co_generator:
canv, i, j, k = my_co.plot_ellipse('x',
'px',
'mm',
'MeV/c',
canvas=canv)
except ValueError:
pass
test_co = [self.co[key] for key in my_co.keys]
for i, item in enumerate(my_co.centre):
self.assertAlmostEqual(my_co.centre[i], test_co[i], 3)
def test_init(self):
"""test EllipseClosedOrbitFinder.__init__"""
co_finder = EllipseClosedOrbitFinder(self.tracking, self.co)
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]