def generate_new_sample_state(self):
"""
Return a list of "particle group"
Each group is a tuple (positions, orientations)
where the positions and orientations have one line
per particle in the sample at this state.
"""
particle_groups = []
particle_distribution = np.random.multinomial(self.n_part_per_shot,
self.ratios)
particle_dict = {
self.particles[i]: n_particles
for i, n_particles in enumerate(particle_distribution)
}
part_states, part_positions = distribute_particles(
particle_dict,
self.beam.get_focus()[0] / 2, self.jet_radius, self.gamma)
part_states = np.array(part_states)
for i in range(self.n_particle_kinds):
n_particles = particle_distribution[i]
orientations = psg.get_random_quat(n_particles)
positions = part_positions[part_states == self.particles[i]]
particle_groups.append((positions, orientations))
return particle_groups
def get_ref_orientation(self):
"""
Return the orientation.
"""
if self._ref_orientation is None:
return psg.get_random_quat(1)
ref_orientation = self._ref_orientation
return ref_orientation
def get_next_orientation(self):
"""
Return the next orientation.
"""
if self._orientations is None:
return psg.get_random_quat(1)
if self._i_orientations >= len(self._orientations):
raise StopIteration("No more orientation available.")
orientation = self._orientations[self._i_orientations, None]
self._i_orientations += 1
return orientation
def get_next_part_orientation(self):
"""
Return the next orientation.
"""
if self._part_orientations is None:
return psg.get_random_quat(1)
if self._i_part_orientations >= len(self._part_orientations):
raise StopIteration("No more orientation available.")
if self.multi_particle_hit:
part_orientation = self._part_orientations[
self._i_part_orientations]
else:
part_orientation = self._part_orientations[
self._i_part_orientations, None]
self._i_part_orientations += 1
return part_orientation
def generate_new_sample_state(self):
"""
Return a list of "particle group"
Each group is a tuple (positions, orientations)
where the positions and orientations have one line
per particle in the sample at this state.
"""
particle_groups = []
if self.n_ref_per_shot >= self.n_part_per_shot:
# reference cluster
reference_distribution = np.random.multinomial(
self.n_ref_per_shot, self.ref_ratios)
ref_dict = {
self.reference[i]: n_reference
for i, n_reference in enumerate(reference_distribution)
}
ref_states, ref_positions = distribute_particles(
ref_dict,
self.beam.get_focus()[0] / 2,
jet_radius=1e-4,
gamma=1.)
ref_states = np.array(ref_states)
for i in range(self.n_reference_kinds):
n_reference = reference_distribution[i]
ref_orientations = psg.get_random_quat(n_reference)
ref_positions = ref_positions[ref_states == self.reference[i]]
# particle cluster
particle_distribution = np.random.multinomial(
self.n_part_per_shot, self.part_ratios)
part_dict = {
self.particles[i]: n_particles
for i, n_particles in enumerate(particle_distribution)
}
part_states, part_positions = distribute_particles(
ref_dict,
self.beam.get_focus()[0] / 2,
jet_radius=1e-4,
gamma=1.)
part_positions[:, 1] = ref_positions[:part_positions.shape[0], 1]
part_states = np.array(part_states)
for i in range(self.n_particle_kinds):
n_particles = particle_distribution[i]
part_orientations = psg.get_random_quat(n_particles)
part_positions = part_positions[part_states ==
self.particles[i]]
else:
# particle cluster
particle_distribution = np.random.multinomial(
self.n_part_per_shot, self.part_ratios)
part_dict = {
self.particles[i]: n_particles
for i, n_particles in enumerate(particle_distribution)
}
part_states, part_positions = distribute_particles(
ref_dict,
self.beam.get_focus()[0] / 2,
jet_radius=1e-4,
gamma=1.)
part_states = np.array(part_states)
for i in range(self.n_particle_kinds):
n_particles = particle_distribution[i]
part_orientations = psg.get_random_quat(n_particles)
part_positions = part_positions[part_states ==
self.particles[i]]
# reference cluster
reference_distribution = np.random.multinomial(
self.n_ref_per_shot, self.ref_ratios)
ref_dict = {
self.reference[i]: n_reference
for i, n_reference in enumerate(reference_distribution)
}
ref_states, reference_positions = distribute_particles(
ref_dict,
self.beam.get_focus()[0] / 2,
jet_radius=1e-4,
gamma=1.)
ref_positions[:, 1] = part_positions[:ref_positions.shape[0], 1]
ref_states = np.array(reference_states)
for i in range(self.n_reference_kinds):
n_reference = reference_distribution[i]
ref_orientations = psg.get_random_quat(n_reference)
ref_positions = ref_positions[ref_states == self.reference[i]]
positions = np.concatenate((ref_positions, part_positions), axis=0)
orientations = np.concatenate((ref_orientations, part_orientations),
axis=0)
particle_groups.append((positions, orientations))
return particle_groups