def setUp(self):
envelope_tool.set_share_dirs()
self.red_dict = {
'pid': 13,
'px': 1.,
'py': 2.,
'pz': 3.,
'x': -1,
'y': -2,
'z': -3
}
self.hit_red = Hit.new_from_dict(self.red_dict)
self.hit_full = Hit.new_from_dict(
{
'pid': 13,
'px': 1.,
'py': 2.,
'pz': 3.,
'x': -1,
'y': -2,
'z': -3,
'mass': common.pdg_pid_to_mass[13],
'charge': common.pdg_pid_to_charge[13]
}, 'energy')
self.main_window = envelope_tool.MainWindow()
self.main_window.window.get_frame("&Beam Setup", "button").Clicked()
self.beam_setup = self.main_window.beam_setup
self.cov = CovarianceMatrix()
for i in range(1, 7):
for j in range(1, 7):
self.cov.set_element(i, j, i * j)
self.cov.set_element(i, i, i * i + 50.)
def test_track_many(self):
"""
Test MatrixTracking.track_many
"""
tracking = MatrixTracking(self.matrix_list, self.offset_list,
self.offset_in)
hit_in = Hit.new_from_dict(
{
"mass": Common.pdg_pid_to_mass[2212],
"charge": 1.,
"pid": 2212,
"pz": 1.,
"x": 1.,
"py": 1.,
"energy": 1001.
}, "pz")
hit_in_2 = Hit.new_from_dict(
{
"mass": Common.pdg_pid_to_mass[2212],
"charge": 1.,
"pid": 2212,
"pz": 1.,
"x": 1.,
"py": 1.,
"energy": 1002.
}, "pz")
hit_list_of_lists = tracking.track_many([hit_in] * 3 + [hit_in_2])
for hit_list in hit_list_of_lists[1:-1]:
self.assertEqual(hit_list, hit_list_of_lists[0])
self.assertNotEqual(hit_list_of_lists[-1], hit_list_of_lists[0])
self.assertEqual(tracking.last, hit_list_of_lists)
def plot_decoupled_da(self, row_data, da_key, index, axis1, axis2,
max_n_points, variables, acceptance):
hit_data = row_data[da_key]
da_row = self.get_da_row(hit_data, max_n_points)
tm = row_data["ref_track"][0]
tm = copy.deepcopy(row_data["tm"])
for i, row in enumerate(tm):
tm[i] = row[1:5]
tm = DecoupledTransferMatrix(tm)
name = da_key + "_" + axis1 + " vs " + axis2
u_data, v_data, au_data, av_data, a4_data = [], [], [], [], []
closed_orbit = Hit.new_from_dict(row_data["ref_track"][0])
for i, hit_list in enumerate(hit_data):
hit_list = [
Hit.new_from_dict(hit_dict) for hit_dict in hit_list[1]
]
seed_list = [self.get_seed(hit, closed_orbit) for hit in hit_list]
seed_list = [tm.decoupled(seed) for seed in seed_list]
u = [seed[0] for seed in seed_list]
pu = [seed[1] for seed in seed_list]
v = [seed[2] for seed in seed_list]
pv = [seed[3] for seed in seed_list]
u_data.append((u, pu))
v_data.append((v, pv))
#turns = range(len(seed_list))
#amp_data_u = [self.get_amplitude(tm, seed, [0, 1]) for seed in seed_list]
##amp_data_v = [self.get_amplitude(tm, seed, [2, 3]) for seed in seed_list]
#amp_data_4d = [self.get_amplitude(tm, seed, [0, 1, 2, 3]) for seed in seed_list]
self.make_plot(u_data, da_row, da_key + " " + str(index) + ": u vs pu",
"u", "pu")
self.make_plot(v_data, da_row, da_key + " " + str(index) + ": v vs pv",
"v", "pv")
def get_toy_fets_ring(self):
"""Parameters version 0.5 N=15 FETS ring"""
self.toy = ToyLongitudinal()
self.start = Hit.new_from_dict(self.co_list[0]['hits'][0])
self.end = Hit.new_from_dict(self.co_list[-1]['hits'][0])
self.toy.r0 = self.start['r']
self.toy.p0 = self.start['p']
self.toy.k = 7.1
self.toy.n_stations = 4
self.toy.v_eff = 5e-3
self.toy.phi_s = 0.1
self.toy.n_turns = 2000
self.toy.p1 = self.end['p']
self.toy.scallop_factor = 1073.47 / self.toy.tof(self.toy.p0)
self.toy.omega = 2. * math.pi * 0.001 / self.toy.tof(self.toy.p0)
self.toy.cavity_phase = [0., 0.25, 0.5,
0.75] # fixed actual phase of cavity
self.toy.azimuthal_angle = [0.] + [0., 0.25, 0.5, 0.75
] + [1.] # position of cavity
self.toy.setup_lookup(self.co_file, n_cells=15)
print("Built ring with frequency", self.toy.omega/2./math.pi, \
"scallop factor", self.toy.scallop_factor)
self.toy.plot_dir = self.plot_dir
self.toy.output_dir = self.tmp_dir
self.toy.print_radius()
self.toy.setup_phi_eff()
canvas = None
for station, color in [(0, ROOT.kRed), (1, ROOT.kGreen),
(2, ROOT.kBlue), (3, ROOT.kMagenta)]:
canvas, hist, graph = self.toy.plot_radial_dependence(
color, station, canvas)
def get_co(self, item):
var_list = ["x", "x'", "y", "y'"]
units = {"x":1., "x'":1e3, "y":1., "y'":1e3}
print("Ref track length", len(item["ref_track"]), item["ref_track"][0])
print(item.keys())
ref_0 = Hit.new_from_dict(item["ref_track"][0])
ref_1 = Hit.new_from_dict(item["ref_track"][2])
delta = [(ref_0[var]-ref_1[var])*units[var] for var in var_list]
delta_mag = sum([x*x for x in delta])**0.5
return delta_mag
def hit_global_weight_test():
Hit.clear_global_weights()
hit = Hit.new_from_dict({
"spill": 1,
"event_number": 2,
"particle_number": 3
})
if (abs(hit.get("global_weight") - 1.) > 1e-9): return 'fail'
hit.set('global_weight', 0.5)
if (abs(hit.get_global_weight() - 0.5) > 1e-9): return 'fail'
Hitcore.clear_global_weights()
if (abs(hit.get_global_weight() - 1.) > 1e-9): return 'fail'
return 'pass'
def get_new_seed(self, config_seed):
if len(self.output_list) == 0 or len(self.output_list[-1]) == 0:
return config_seed
elif len(self.output_list[-1]) == 1:
seed = Hit.new_from_dict(self.output_list[-1][-1]["seed_hit"])
return [seed[var] for var in self.var_list]
else:
seed0 = Hit.new_from_dict(self.output_list[-1][-2]["seed_hit"])
seed1 = Hit.new_from_dict(self.output_list[-1][-1]["seed_hit"])
s0 = [seed0[var] for var in self.var_list]
s1 = [seed1[var] for var in self.var_list]
s2 = [2 * s1[i] - s0[i] for i in range(4)]
return s2
def load_tracks(self, track_vector):
loaded_track_vector = []
for track in track_vector:
if abs(track.GetParticleId()
) == 211 and track.GetKillReason() != "":
print "Pion killed because", track.GetKillReason()
hit = {}
hit["x"] = track.GetInitialPosition().x()
hit["y"] = track.GetInitialPosition().y()
hit["z"] = track.GetInitialPosition().z()
hit["px"] = track.GetInitialMomentum().x()
hit["py"] = track.GetInitialMomentum().y()
hit["pz"] = track.GetInitialMomentum().z()
hit["pid"] = track.GetParticleId()
try:
hit["mass"] = xboa.common.pdg_pid_to_mass[abs(hit["pid"])]
except KeyError:
hit["mass"] = 0.
try:
hit["charge"] = xboa.common.pdg_pid_to_charge[hit["charge"]]
except KeyError:
hit["charge"] = 0.
hit = {}
hit["x"] = track.GetFinalPosition().x()
hit["y"] = track.GetFinalPosition().y()
hit["z"] = track.GetFinalPosition().z()
hit["px"] = track.GetFinalMomentum().x()
hit["py"] = track.GetFinalMomentum().y()
hit["pz"] = track.GetFinalMomentum().z()
hit["pid"] = track.GetParticleId()
try:
hit["mass"] = xboa.common.pdg_pid_to_mass[abs(hit["pid"])]
except KeyError:
hit["mass"] = 0.
try:
hit["charge"] = xboa.common.pdg_pid_to_charge[hit["charge"]]
except KeyError:
hit["charge"] = 0.
loaded_track_initial = {
"hit": Hit.new_from_dict(hit, "energy"),
"detector": "mc_track_initial"
}
loaded_track_final = {
"hit": Hit.new_from_dict(hit, "energy"),
"detector": "mc_track_final",
}
loaded_track_vector += [loaded_track_initial, loaded_track_final]
return {}, loaded_track_vector
def _test_content_hypersphere(self):
# disabled because it is slow
print("Test content hypersphere")
for i in range(5):
numpy.seterr(all='ignore')
test_bunch = Bunch()
content_predicted = numpy.linalg.det(self.diag)*math.pi**2/2.
test_bunch = Bunch.new_hit_shell(7,
4.*self.diag,
['x', 'y', 'px', 'py'],
'')
test_bunch.append(Hit.new_from_dict({'x':0., 'y':0., 'px':0., 'py':0.}))
for hit in test_bunch:
hit['x'] += 1.*i
hit['px'] += 2.*i
hit['y'] += 3.*i
hit['py'] += 4.*i
my_weights = VoronoiWeighting(['x', 'y', 'px', 'py'],
numpy.array([[2., 0., 0., 0.],
[0., 1., 0., 0.],
[0., 0., 1., 0.],
[0., 0., 0., 1.],]))
my_weights.apply_weights(test_bunch, False)
my_weights.plot_two_d_projection(['x', 'y'])
my_weights.plot_two_d_projection(['px', 'py'])
#self.assertEqual(len(my_weights.tile_content_list), len(test_bunch)+1)
content_actual = sum(my_weights.tile_content_list)
print("Content", content_actual, "Predicted", content_predicted)
self.assertAlmostEqual(content_actual, content_predicted, 0)
input()
def get_tracking_object():
"""
Create a tracking object. This tracking object uses a list of matrices to
generate tracking. The tracking goes like
- x is a conventional phase space ellipse
- y has phase advance of exactly pi
- t has phase advance is complex i.e. on a resonance.
Each cell or turn we advance through another transfer matrix. The closed
orbit is generated by a hard coded offset.
Return value is a tuple consisting of a hit object on the closed_orbit and
a tracking object that will propagate hit objects through the matrices.
"""
number_of_turns = 5
# x is conventional ellipse
# y is phase advance = pi
# t is phase advance = complex
data = [[1.0, 0.5, 0.0, 0.0, 0.0, 0.0], [-0.5, 0.75, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, -1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 2.0, 2.0], [0.0, 0.0, 0.0, 0.0, 2.0, 2.5]]
matrix = numpy.matrix(data)
matrix_list = [matrix**(i + 1) for i in range(number_of_turns)]
offset = numpy.matrix([10., 7., 0., 0., 0., 1000.])
offset_list = [offset] * number_of_turns
offset_in = numpy.matrix([10., 7., 0., 0., 0., 1000.])
tracking = MatrixTracking(matrix_list, offset_list, offset_in)
closed_orbit_hit = Hit.new_from_dict(
{
'x': 10.,
'px': 7.,
'energy': 1000.,
'pid': 2212,
'mass': common.pdg_pid_to_mass[2212]
}, 'pz')
return closed_orbit_hit, tracking
def generate_h5_step(self, h5_file, station):
for key in h5_file.keys():
if key[:5] != "Step#":
if self.verbose > 10:
print("Skipping", key)
continue
n_steps = len(h5_file[key]["x"])
if self.verbose > 10:
print("Found", key, "in", n_steps, "events in", h5_file.filename)
h5_step = h5_file[key]
for i in range(n_steps):
hit_dict = {}
for key in "pid", "mass", "charge":
hit_dict[key] = self.ref[key]
for h5_key, xboa_key in self.h5_key_to_xboa_key.items():
hit_dict[xboa_key] = h5_step[h5_key][i]
if xboa_key in self.units:
hit_dict[xboa_key] *= self.units[xboa_key]
for key in ["px", "py", "pz"]:
hit_dict[key] *= self.ref["mass"]
hit_dict["station"] = station
hit = Hit.new_from_dict(hit_dict, "energy")
if self.verbose > 10:
print(" h5", format(hit["event_number"], "4d"), end=' ')
for key in self.print_keys:
print(str(round(hit[key], 3)).ljust(8), end=' ')
print()
yield hit["event_number"], hit
def test_track_many(self):
"""
Test MAUSTracking.track_many
"""
tracking = MAUSTracking(self.config)
hit_in = Hit.new_from_dict(
{
"mass": Common.pdg_pid_to_mass[13],
"charge": 1.,
"energy": 226.,
"pid": -13,
"pz": 200.
}, "energy")
hit_list_of_lists = tracking.track_many([hit_in] * 5)
self.assertEqual(len(hit_list_of_lists), 5)
for i, hit_list in enumerate(hit_list_of_lists):
for j, hit_out in enumerate(hit_list):
for key in [
"pid", "mass", "charge", "x", "y", "energy", "px",
"py", "pz"
]:
mess = "Failed on "+key+" "+\
str(hit_in[key])+", "+str(hit_out[key])
self.assertAlmostEqual(hit_in[key],
hit_out[key],
3,
msg=mess)
self.assertEqual(hit_out["z"], j * 1000.)
self.assertEqual(hit_out["event_number"], i)
self.assertEqual(hit_out["particle_number"], 1)
self.assertEqual(hit_out["spill"], 0)
self.assertEqual(tracking.last, hit_list_of_lists)
def setUp(self):
"""Setup the EllipseClosedOrbitFinder"""
self.num_turns = 5
# x is conventional ellipse
# y is phase advance = pi
# t is phase advance = complex
data = [[1.0, 0.5, 0.0, 0.0, 0.0,
0.0], [-0.5, 0.75, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, -1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 2.0, 2.0],
[0.0, 0.0, 0.0, 0.0, 2.0, 2.5]]
matrix = numpy.matrix(data)
self.matrix_list = [matrix**(i + 1) for i in range(self.num_turns)]
offset = numpy.matrix([10., 7., 0., 0., 0., 1000.])
self.offset_list = [offset] * self.num_turns
self.offset_in = numpy.matrix([10., 7., 0., 0., 0., 1000.])
self.tracking = MatrixTracking(self.matrix_list, self.offset_list,
self.offset_in)
self.co = Hit.new_from_dict({
'x': 10.,
'px': 7.,
'energy': 1000.,
'pid': 2212,
'mass': common.pdg_pid_to_mass[2212]
})
if abs(numpy.linalg.det(matrix) - 1.0) > 1e-9:
raise ValueError("TM determinant should be 1, got "+\
str(numpy.linalg.det(matrix)))
def make_hit_list(self, deltas):
defaults = {
'pz': self.pz,
'mass': self.m_mu,
'charge': 1.,
'pid': -13,
'z': self.z_start
}
hit_list = [Hit.new_from_dict(defaults, 'energy')]
delta_keys = ['x', 'y', 'px', 'py']
for parity in [-1., +1.]:
for i in range(4):
defaults[delta_keys[i]] = parity * deltas[i]
hit_list.append(Hit.new_from_dict(defaults, 'energy'))
defaults[delta_keys[i]] = 0.
return hit_list
def load_tof_sp(self, tof_sp, station):
xerr = tof_sp.GetGlobalPosXErr()
yerr = tof_sp.GetGlobalPosYErr()
cov = [
[0.0025, 0., 0., 0., 0., 0.],
[0., xerr, 0., 0., 0., 0.],
[0., 0., yerr, 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.],
]
sp_dict = {
"x": tof_sp.GetGlobalPosX(),
"y": tof_sp.GetGlobalPosY(),
"z": tof_sp.GetGlobalPosZ(),
"t": tof_sp.GetTime() + self.time_offsets[station],
}
loaded_sp = {
"hit": Hit.new_from_dict(sp_dict),
"detector": station,
"covariance": cov,
"dt": tof_sp.GetDt(),
"charge_product": tof_sp.GetChargeProduct(),
"charge": tof_sp.GetCharge(),
}
return loaded_sp
def test_track_one(self):
"""
Test MAUSTracking.track_one
"""
# BUG: G4 does not reinitialise properly if run after test_tracking_init
tracking = MAUSTracking(self.config, 99)
hit_in = Hit.new_from_dict(
{
"mass": Common.pdg_pid_to_mass[2212],
"charge": 1.,
"energy": 10000.,
"pid": 2212,
"pz": 1.
}, "pz")
hit_list = tracking.track_one(hit_in)
self.assertEqual(len(hit_list), 6)
for i, hit_out in enumerate(hit_list):
for key in [
"pid", "mass", "charge", "x", "y", "energy", "px", "py",
"pz"
]:
mess = "Failed on "+key+" "+\
str(hit_in[key])+", "+str(hit_out[key])
self.assertAlmostEqual(hit_in[key], hit_out[key], 3, msg=mess)
self.assertEqual(hit_out["z"], i * 1000.)
self.assertEqual(hit_out["event_number"], 0)
self.assertEqual(hit_out["particle_number"], 1)
self.assertEqual(hit_out["spill"], 0)
hit_list = tracking.track_one(hit_in)
self.assertEqual(len(hit_list), 6)
for i, hit_out in enumerate(hit_list):
self.assertEqual(hit_out["event_number"], 0)
self.assertEqual(hit_out["particle_number"], 1)
self.assertEqual(hit_out["spill"], 1)
self.assertEqual(tracking.last, [hit_list])
def read_one_line(self, line, station):
words = line.split()
hit_dict = {}
for key in "pid", "mass", "charge":
hit_dict[key] = self.ref[key]
for i, key in enumerate(["x", "z", "y"]):
hit_dict[key] = float(words[i + 1])
for i, key in enumerate(["px", "pz", "py"]):
hit_dict[key] = float(words[i + 4]) * self.ref["mass"]
event = int(words[7])
hit_dict["event_number"] = int(words[7])
hit_dict["station"] = station
x = hit_dict["x"]
y = hit_dict["z"]
px = hit_dict["px"]
py = hit_dict["pz"]
phi = math.atan2(y, x)
hit_dict["t"] = float(words[9])
hit_dict["x"] = +x * math.cos(phi) + y * math.sin(phi)
hit_dict["z"] = -x * math.sin(phi) + y * math.cos(phi)
hit_dict["px"] = +px * math.cos(phi) + py * math.sin(phi)
# go through file line by line reading hit data
hit_dict["pz"] = -px * math.sin(phi) + py * math.cos(phi)
hit = Hit.new_from_dict(hit_dict, "energy")
return event, hit
def calculate_da(row_data, da_key, n_points):
axis1, axis2 = {"x_da": ("x", "x'"), "y_da": ("y", "y'")}[da_key]
points = []
mass = common.pdg_pid_to_mass[2212]
hit_list = []
for i, tmp_list in enumerate(row_data[da_key]):
print(len(tmp_list), tmp_list[0])
# hit_list[0] is the x/y offset in mm; hit_list[1] is the corresponding
# set of probe hits
if len(tmp_list[1]) < n_points:
break
hit_list = tmp_list
hit_list = [Hit.new_from_dict(hit_dict) for hit_dict in hit_list[1]]
points = numpy.array([[hit[axis1], hit[axis2]] for hit in hit_list])
x_list = [hit[axis1] for hit in hit_list]
print(min(x_list), max(x_list))
y_list = [hit[axis2] for hit in hit_list]
print(min(y_list), max(y_list))
pz = hit_list[0]['pz']
print(math.pi * (max(x_list) - min(x_list)) * (max(y_list) - min(y_list)) /
4. * mass / pz)
if len(points):
geometric_acceptance = get_area(points) / math.pi
# beta gamma * geometric emittance
normalised_acceptance = geometric_acceptance * pz / mass
return normalised_acceptance
else:
return 0.
def test_make_hit(self):
format_list = ['', 'x', 'y', 'z', 'energy', 'pid']
format_units_dict = {
'x': 'mm',
'y': 'm',
'z': 'cm',
'energy': '',
'pz': '',
'pid': ''
}
filehandle = open(tmp_name(), 'w')
if not filehandle:
raise unittest.SkipTest("Failed to open filehandle")
filehandle.write('-99. 1. 2. 3. 4. -13\n')
filehandle.write('-99. 1. 2. 3. 4. -13\n')
filehandle.write('-99. 1. 2. 3. 4.')
filehandle.close()
ref_hit = Hit.new_from_dict(
{
'x': 1.,
'y': 2000.,
'z': 30.,
'energy': 4.,
'pid': -13,
'charge': 1.,
'mass': common.pdg_pid_to_mass[13]
}, 'pz')
filehandle = open(tmp_name(), 'r')
fac = UserHitFactory(format_list, format_units_dict, filehandle, 'pz')
test_hit = fac.make_hit()
self.assertEqual(test_hit, ref_hit)
test_hit = fac.make_hit()
self.assertEqual(test_hit, ref_hit)
self.assertRaises(IOError, fac.make_hit)
filehandle.close()
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 setup_lookup(self, filename, n_cells):
json_in = [json.loads(line) for line in open(filename).readlines()]
my_list = []
for json_line in json_in:
hits = [Hit.new_from_dict(item) for item in json_line['hits']]
r = numpy.mean([hit['r'] for hit in hits])
p = numpy.mean([hit['p'] for hit in hits])
tof_list = [
hits[i + 1]['t'] - hits[i]['t']
for i, hit in enumerate(hits[1:])
]
mean_tof = numpy.mean(tof_list)
tof_list = [
tof for tof in tof_list if abs(tof - mean_tof) < mean_tof / 2.
]
mean_tof = numpy.mean(tof_list)
print("Mean tof", mean_tof, "rms tof", numpy.std(tof_list))
t = mean_tof * n_cells
my_list.append((r, p, t))
my_list = sorted(my_list)
r_list, p_list, t_list = list(zip(*my_list))
for i, p in enumerate(p_list):
print(i, format(p, "8.4g"), format(r_list[i], "8.4g"),
format(t_list[i], "8.4g"))
self.lookup_tof = interp1d(p_list, t_list, kind='cubic')
self.lookup_tof.bounds_error = False
self.lookup_tof.fill_value = 'extrapolate'
self.lookup_radius = interp1d(p_list, r_list, kind='cubic')
self.lookup_radius.bounds_error = False
self.lookup_radius.fill_value = 'extrapolate'
def load_closed_orbits(self):
for i, ref_subs in enumerate(self.config.substitution_list):
fin = open(self.config.run_control["output_dir"] +
"/find_rf_closed_orbits/find_rf_time_" + str(i) +
".out")
co_list = [json.loads(line) for line in fin.readlines()]
et_list = []
self.closed_orbit_list.append(
[Hit.new_from_dict(item["hits"][0]) for item in co_list])
for co in co_list:
energy = co["hits"][0]["energy"] - self.mass
hit_times = [hit["t"] for hit in co["hits"][1:]]
delta_times = [
t1 - hit_times[i] for i, t1 in enumerate(hit_times[1:])
]
mean_time = numpy.mean(delta_times) #]
delta_times = [
time for time in delta_times
if abs(time - mean_time) < mean_time / 2.
]
print("One Cell Time", mean_time, "Sigma",
numpy.std(delta_times))
mean_time = numpy.mean(
delta_times) * self.config.find_rf_parameters["n_cells"]
et_list.append((energy, mean_time))
self.energy_time_list.append(et_list)
def load_scifi_space_points(self, scifi_event):
space_points_out = []
for space_point in scifi_event.spacepoints():
# if we require triplets and space point is not a triplet, skip it
if self.config.will_require_triplets and \
space_point.get_channels().GetEntries() != 3:
continue
position = space_point.get_global_position()
sp_dict = {
"x": position.x(),
"y": position.y(),
"z": position.z(),
}
space_points_out.append({
"hit":
Hit.new_from_dict(sp_dict),
"detector": ["tku_sp_", "tkd_sp_"][space_point.get_tracker()],
"n_channels":
space_point.get_channels().GetEntries(),
"is_used":
space_point.is_used(),
"channels":
[chan.get_channel() for chan in space_point.get_channels()],
"npe": [chan.get_npe() for chan in space_point.get_channels()],
})
space_points_out[-1]["detector"] += str(space_point.get_station())
return space_points_out
def load_json(file_name):
print "Loading", file_name
fin = open(file_name)
data = [json.loads(line) for line in fin.readlines()]
for item in data:
item["seed"] = [Hit.new_from_dict(a_dict) for a_dict in item["seed"]]
print "Loaded"
return data
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(self):
offset_in = numpy.matrix(numpy.zeros([1, 6]))
offset_out = [numpy.matrix(numpy.zeros([1, 6]))]
matrices = [numpy.matrix(numpy.zeros([6, 6]))]
for i in range(6):
matrices[0][i, i] = i
self.tracking = TimeoutTracking(matrices, offset_out, offset_in, 0)
self.hit_list = [Hit.new_from_dict({'x':i}) for i in range(1000)]
def load_json(file_name):
print "Loading", file_name
fin = open(file_name)
data = [json.loads(line) for line in fin.readlines()]
for item in data:
item["seed"] = [Hit.new_from_dict(a_dict) for a_dict in item["seed"]]
print "Loaded"
return data
def plot_one_phi(row_data, da_key, axis1, axis2, max_n_points, variables):
global ROOT_OBJECTS
hit_data = row_data[da_key]
name = "phi_" + da_key + "_" + axis1 + " vs " + axis2
title = get_title(variables)
name += " " + title
canvas = None
tune_data = []
area_data = []
da_row = get_da_row(hit_data, max_n_points)
print("Plot on phi")
cholesky_canvas = common.make_root_canvas(name + "_cholesky")
cholesky_canvas.SetCanvasSize(1000, 1000)
delta = (da_row + 1) * 0.4 + 1.
hist = common.make_root_histogram("", [-100.],
"u",
100, [-100.],
"u'",
100,
xmin=-delta,
xmax=+delta,
ymin=-delta,
ymax=+delta)
hist.Draw()
for i, hit_list in enumerate(hit_data[:da_row + 1]):
hit_list = [Hit.new_from_dict(hit_dict) for hit_dict in hit_list[1]]
finder = DPhiTuneFinder()
finder.u = [hit[axis1] for hit in hit_list[1:]]
finder.up = [hit[axis2] for hit in hit_list[1:]]
area_src = numpy.array([[hit[axis1], hit[axis2]] for hit in hit_list])
try:
finder.get_tune()
an_area = get_area(area_src)
except Exception:
print("Failed to unpack data for phi da plot")
continue
cholesky_canvas, hist, graph = finder.plot_cholesky_space(
cholesky_canvas, 1. + i * 0.2)
if i < da_row:
color = ROOT.TColor(10000 + len(ROOT_OBJECTS), 0.,
1. - i * 1. / da_row, 0.)
ROOT_OBJECTS.append(color)
graph.SetMarkerColor(ROOT.kGreen)
ROOT_OBJECTS += [hist, graph]
tune_data += copy.deepcopy(finder.dphi)
area_data += [an_area for i in finder.dphi]
print(" Area", an_area, end=' ')
print(" tune", numpy.mean(tune_data), "+/-", numpy.std(tune_data))
canvas = common.make_root_canvas(name)
canvas.Draw()
hist, graph = common.make_root_graph(name, area_data, "Amplitude [mm]",
tune_data, "Fractional tune")
ROOT_OBJECTS += [hist, graph]
hist.SetTitle(da_key)
hist.Draw()
graph.SetMarkerStyle(24)
graph.Draw("p same")
return canvas, cholesky_canvas
def setUp(self):
envelope_tool.set_share_dirs()
self.red_dict = {'pid':13, 'px':1., 'py':2., 'pz':3.,
'x':-1, 'y':-2, 'z':-3}
self.hit_red = Hit.new_from_dict(self.red_dict)
self.hit_full = Hit.new_from_dict({'pid':13, 'px':1., 'py':2., 'pz':3.,
'x':-1, 'y':-2, 'z':-3,
'mass':common.pdg_pid_to_mass[13],
'charge':common.pdg_pid_to_charge[13]},
'energy')
self.main_window = envelope_tool.MainWindow()
self.main_window.window.get_frame("&Beam Setup", "button").Clicked()
self.beam_setup = self.main_window.beam_setup
self.cov = CovarianceMatrix()
for i in range(1, 7):
for j in range(1, 7):
self.cov.set_element(i, j, i*j)
self.cov.set_element(i, i, i*i+50.)
def _reference(self, energy):
"""Generate a reference particle"""
hit_dict = {}
hit_dict["pid"] = 2212
hit_dict["mass"] = common.pdg_pid_to_mass[2212]
hit_dict["charge"] = 1
hit_dict["x"] = 4600.
hit_dict["kinetic_energy"] = energy
return Hit.new_from_dict(hit_dict, "pz")
def write(self):
hit_list = []
for item in self.tgt_dist:
my_dict = {"mass": self.mu_mass, "pid": -13}
for i, key in enumerate(self.keys):
my_dict[key] = item[i]
hit_list.append(Hit.new_from_dict(my_dict, "energy"))
bunch = Bunch.new_from_hits(hit_list)
bunch.hit_write_builtin(self.format, self.output_file)
print "Writing bunch of length", len(bunch)
def setUp(self):
offset_in = numpy.matrix(numpy.zeros([1, 6]))
offset_out = [numpy.matrix(numpy.zeros([1, 6]))]
matrices = [numpy.matrix(numpy.zeros([6, 6]))]
for i in range(6):
matrices[0][i, i] = i
self.tracking = TimeoutTracking(matrices, offset_out, offset_in, 0)
self.hit_list = [Hit.new_from_dict({'x':1, 'y':2})]
self.tracking_process = ThisNodeProcess(self.tracking)
self.time_constant = 10
def reference(energy):
"""
Generate a reference particle
"""
hit_dict = {}
hit_dict["pid"] = CONFIG.tracking["pdg_pid"]
hit_dict["mass"] = common.pdg_pid_to_mass[abs(hit_dict["pid"])]
hit_dict["charge"] = 1
hit_dict["x"] = 0.
hit_dict["kinetic_energy"] = energy
return Hit.new_from_dict(hit_dict, "pz")
def get_hit(x, y, z, t, px, py, pz):
hit_dict = {
'mass':common.pdg_pid_to_mass[13],
'pid':-13,
'x':x,
'y':y,
'z':z,
'px':px,
'py':py,
'pz':pz,
't':0.
}
hit = Hit.new_from_dict(hit_dict, 'energy')
return hit
def test_beam_setup_cancel(self):
"""Test beam_setup cancel button"""
test_matrix = self.beam_setup.get_matrix()
test_ref = self.beam_setup.get_reference()
self.beam_setup.window.get_frame("&Okay", "button").Clicked()
self.main_window.window.get_frame("&Beam Setup", "button").Clicked()
self.beam_setup = self.main_window.beam_setup
self.beam_setup.set_matrix(self.cov)
self.beam_setup.set_reference(self.hit_red)
self.beam_setup.window.get_frame("&Cancel", "button").Clicked()
for var in Hit.get_variables():
self.assertAlmostEqual(self.main_window.lattice.ref_list[0][var],
test_ref[var], 3)
self._ell_equal(self.main_window.lattice.ellipse_list[0], test_matrix)
def process(self, json_spill_doc):
"""
Generate primary particles for a spill.
@param json_spill_doc unicode string containing the json spill
- In counter mode, iterates over each beam and samples required number
of particles from the parent distribution.
- In overwrite_existing or binomial mode, randomly samples from each of
the available beam distributions according to the relative weight of
each beam.
@returns a string with the json spill
"""
spill = {}
try:
spill = json.loads(json_spill_doc)
spill = self.__process_check_spill(spill)
new_particles = self.__process_gen_empty(spill)
for index, particle in enumerate(new_particles):
# if beam IO, then read hits from file and fill spill
if (self.use_beam_file):
spill_hit = Hit.new_from_read_builtin(self.beam_file_format,
self.bm_fh)
primary = spill_hit.get_maus_dict('maus_json_primary')[0]
self.beam_seed = self.beam_seed + 1
primary["random_seed"] = self.beam_seed
primary['position']['x'] += \
self.beam_file_insert_position['x']
primary['position']['y'] += \
self.beam_file_insert_position['y']
primary['position']['z'] += \
self.beam_file_insert_position['z']
particle["primary"] = primary
else:
a_beam = self.__process_choose_beam(index)
particle["primary"] = a_beam.make_one_primary()
if "spin" not in particle["primary"]:
particle["primary"]["spin"] = {"x":0., "y":0., "z":1.}
except Exception: #pylint: disable=W0703
ErrorHandler.HandleException(spill, self)
return json.dumps(spill)
def test_process_a_p_correlation(self):
"""Test beam.py a-p correlation routine"""
self._beam.beam_matrix = numpy.diag([0.1]*6)
self._beam.beam_mean = numpy.zeros((6))
self._beam.a_p_correlation = {"momentum_variable":"pz", "magnitude":0.5}
ref_pz = 200.
hit1 = Hit.new_from_dict({"pz":ref_pz, "mass":105.658}, "energy")
hit2 = hit1.deepcopy()
self._beam._Beam__process_a_p_correlation(hit2)
self.assertEqual(hit1, hit2)
for matrix in [numpy.diag([0.1]*6), numpy.diag([0.2]*6)]:
self._beam.beam_matrix = matrix
hit2 = hit1.deepcopy()
hit3 = hit1.deepcopy()
hit2['x'] = 2.
hit3['px'] = 2.
self._beam._Beam__process_a_p_correlation(hit2)
self._beam._Beam__process_a_p_correlation(hit3)
self.assertAlmostEqual(hit2['pz'], hit3['pz'])
self.assertAlmostEqual(hit2['pz']-ref_pz, ref_pz*0.5*4.*105.658**2)
def get_reference(self):
"""
Get the reference particle (reading from GUI elements)
Returns an xboa.Hit.Hit object
"""
ref_dict = {}
for var in ["x", "y", "z", "px", "py", "pz"]:
var_dict = self.window.get_frame_dict(var, "named_text_entry")
ref_dict[var] = float(var_dict["text_entry"].text_entry.GetText())
pid_dict = self.window.get_frame_dict("pid", "named_text_entry")
try:
ref_dict["pid"] = int(pid_dict["text_entry"].text_entry.GetText())
ref_dict["mass"] = common.pdg_pid_to_mass[abs(ref_dict["pid"])]
ref_dict["charge"] = common.pdg_pid_to_charge[abs(ref_dict["pid"])]
except KeyError:
raise GuiError("Did not recognise reference particle pid")
try:
hit = Hit.new_from_dict(ref_dict, "energy")
except Exception:
raise GuiError("Failed to generate a reference particle")
return hit
def _read_probes(self):
hit_dict_of_lists = {} # maps event number to a list of hit_dicts
# loop over files in the glob, read events and sort by event number
file_list = glob.glob(self.output_name)
for file_name in file_list:
fin = open(file_name)
fin.readline()
# go through file line by line reading hit data
for line in fin.readlines():
words = line.split()
hit_dict = {}
for key in "pid", "mass", "charge":
hit_dict[key] = self.ref[key]
for i, key in enumerate(["x", "z", "y"]):
hit_dict[key] = float(words[i+1])
for i, key in enumerate(["px", "pz", "py"]):
hit_dict[key] = float(words[i+4])*self.ref["mass"]
event = int(words[7])
hit_dict["event_number"] = int(words[7])
hit_dict["station"] = int(words[8])
hit_dict["t"] = float(words[9])
hit_dict["x"] = (hit_dict["x"]**2.+hit_dict["z"]**2.)**0.5
hit_dict["z"] = 0.
if not event in hit_dict_of_lists:
hit_dict_of_lists[event] = []
hit_dict_of_lists[event].append(Hit.new_from_dict(hit_dict, "energy"))
# convert from a dict of list of hits to a list of list of hits
# one list per event
# each list contains one hit per station
events = sorted(hit_dict_of_lists.keys())
hit_list_of_lists = [hit_dict_of_lists[ev] for ev in events]
# sort by time within each event
for i, hit_list in enumerate(hit_list_of_lists):
hit_list_of_lists[i] = sorted(hit_list, key = lambda hit: hit['t'])
self.last = hit_list_of_lists
return hit_list_of_lists
def setUp(self): # pylint: disable=C0103
ref = Hit.new_from_dict({'pid':13, 'px':1., 'py':2., 'pz':3.,
'x':-1, 'y':-2, 'z':-3,
'mass':Common.pdg_pid_to_mass[13],
'charge':Common.pdg_pid_to_charge[13]},
'energy')
self.ref_list = [ref]*4
self.p_dict = {
'e_t':1.,
'b_t':2.,
'e_l':3.,
'b_l':4.,
'a_t':5.,
'a_l':6.,
'l_t':7.,
'd_x':8.,
'dp_x':9.,
'd_y':10.,
'dp_y':11.,
}
self.penn_ellipses = [
maus_cpp.covariance_matrix.create_from_penn_parameters(
mass = ref['mass'],
momentum = ref['p'],
emittance_t = self.p_dict['e_t']+i*12.,
beta_t = self.p_dict["b_t"]+i*12.,
emittance_l = self.p_dict["e_l"]+i*12.,
beta_l = self.p_dict["b_l"]+i*12.,
alpha_t = self.p_dict["a_t"]+i*12.,
alpha_l = self.p_dict["a_l"]+i*12.,
charge = ref["charge"],
bz = 0.,
ltwiddle = self.p_dict["l_t"]+i*12.,
dispersion_x = self.p_dict["d_x"]+i*12.,
dispersion_prime_x = self.p_dict["dp_x"]+i*12.,
dispersion_y = self.p_dict["d_y"]+i*12.,
dispersion_prime_y = self.p_dict["dp_y"]+i*12.,
) for i in range(3)]
self.t_dict = {
'e_x':1.,
'b_x':2.,
'e_y':3.,
'b_y':4.,
'e_l':5.,
'b_l':6.,
'a_x':7.,
'a_y':8.,
'a_l':9.,
'd_x':10.,
'dp_x':11.,
'd_y':12.,
'dp_y':13.,
}
self.twiss_ellipses = [
maus_cpp.covariance_matrix.create_from_twiss_parameters(
mass = ref['mass'],
momentum = ref['p'],
emittance_x = self.t_dict['e_x']+i*14.,
beta_x = self.t_dict["b_x"]+i*14.,
emittance_y = self.t_dict['e_y']+i*14.,
beta_y = self.t_dict["b_y"]+i*14.,
emittance_l = self.t_dict["e_l"]+i*14.,
beta_l = self.t_dict["b_l"]+i*14.,
alpha_x = self.t_dict["a_x"]+i*14.,
alpha_y = self.t_dict["a_y"]+i*14.,
alpha_l = self.t_dict["a_l"]+i*14.,
dispersion_x = self.t_dict["d_x"]+i*14.,
dispersion_prime_x = self.t_dict["dp_x"]+i*14.,
dispersion_y = self.t_dict["d_y"]+i*14.,
dispersion_prime_y = self.t_dict["dp_y"]+i*14.,
) for i in range(3)]
self.field_list = [{
'field_type': 'Multipole',
'position': {'x': 5.0, 'y': 4.0, 'z': 6.0},
'aperture': {'x': 8.0, 'y': 7.0, 'z': 9.0},
'outer': {'x': 10.0, 'y': 11.0, 'z': 12.0},
'scale_factor': -0.5,
'rotation': {'x':math.pi/6., 'y': 0.0, 'z': 0.0},
'field_name': 'D1'
}, {
'field_type': 'Solenoid',
'position': {'x': 5.0, 'y': 4.0, 'z': 6.0},
'aperture': {'x': 8.0, 'y': 7.0, 'z': 9.0},
'outer': {'x': 11.0, 'y': 12.0, 'z': 13.0},
'scale_factor': 0.5,
'rotation': {'x':0.0, 'y': math.pi/6., 'z': 0.0},
'field_name': 'D2'
}]