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 _get_emittance(self, ref, ellipse):
"""Get the beam emittance for given reference hit and ellipse"""
axis_list = self.axis_dict[self.first_var]
el_list = self.el_dict[self.first_var]
bunch = Bunch.new_from_hits([ref])
cov_matrix = [None]*len(el_list)
for cov_i, ell_i in enumerate(el_list):
cov_matrix[cov_i] = [None]*len(el_list)
for cov_j, ell_j in enumerate(el_list):
cov_matrix[cov_i][cov_j] = ellipse.get_element(ell_i, ell_j)
cov_matrix = numpy.array(cov_matrix)
emittance = bunch.get_emittance(axis_list, cov_matrix)
return emittance
def plot_one(self, particles, plot_name):
scraped_list = []
not_scraped_list = []
for i, hit in enumerate(particles):
if self.is_scraped[i]:
scraped_list.append(hit)
else:
not_scraped_list.append(hit)
all_bunch = Bunch.new_from_hits(particles)
scraped_bunch = Bunch.new_from_hits(scraped_list)
not_scraped_bunch = Bunch.new_from_hits(not_scraped_list)
for axes in [('y', 'px'), ('x', 'y')]:
canvas, hist, graph = all_bunch.root_scatter_graph(axes[0],
axes[1],
xmin=-200.,
xmax=200.,
ymin=-200,
ymax=200)
if len(not_scraped_bunch) > 0:
canvas, hist, graph = not_scraped_bunch.root_scatter_graph\
(axes[0], axes[1], canvas=canvas)
graph.SetMarkerStyle(7)
if len(scraped_bunch) > 0:
canvas, hist, graph = scraped_bunch.root_scatter_graph\
(axes[0], axes[1], canvas=canvas)
graph.SetMarkerColor(ROOT.kRed)
graph.SetMarkerStyle(7)
canvas.Update()
canvas.Print(plot_name + "_" + axes[0] + "_vs_" + axes[1] + ".png")
canvas, hist = not_scraped_bunch.root_histogram(axes[0],
"",
axes[1],
"",
xmin=-200.,
xmax=200.,
ymin=-200,
ymax=200)
canvas.Print(plot_name + "_hist_" + axes[0] + "_vs_" + axes[1] +
".png")
def track_bunch(self, bunch):
primary_list = bunch.hits()
hit_list_of_lists = self.maus.track_many(primary_list)
# hit_list_of_lists is sliced by primary; we want to slice by station
n_events = len(primary_list)
# we take the min here to stop
n_stations = min([len(hit_list) for hit_list in hit_list_of_lists])
i_list = range(n_events)
j_list = range(n_stations)
hit_lost_if_losts = [[hit_list_of_lists[i][j] for i in i_list] for j in j_list]
# convert to a bunch
bunch_list = [Bunch.new_from_hits(hit_list) for hit_list in hit_lost_if_losts]
self.bunch_list = bunch_list
return bunch_list
def write(self, z_position = 0., seed = None, output_file = None):
if seed != None:
numpy.random.seed(seed=seed)
if output_file != None:
self.output_file = output_file
hit_list = []
self.tgt_dist = self.kernel.resample(self.n_events).transpose()
for item in self.tgt_dist:
my_dict = {"mass":self.mu_mass, "pid":-13, "z":z_position}
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 track(self, centre_hit, delta):
primary_list = [centre_hit]
for key, value in delta.iteritems():
hit = centre_hit.deepcopy()
hit[key] += value
hit.mass_shell_condition('pz')
primary_list.append(hit)
# qhull needs 8 particles for some reason - so make an extra one
hit = centre_hit.deepcopy()
for key, value in delta.iteritems():
hit[key] += value
hit.mass_shell_condition('pz')
primary_list.append(hit)
bunch = Bunch.new_from_hits(primary_list)
return self.track_bunch(bunch)