def test_add_children(self, particle): """Test the ability to add children to a particle""" event = Event(particle) child1 = Particle(particle_id=Particle.Type.electron, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9, interaction_model=Interaction) child2 = Particle(particle_id=Particle.Type.positron, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9, interaction_model=Interaction) event.add_children(particle, [child1, child2]) assert len(event.roots) == 1 assert len(event) == 3 all_particles = [particle, child1, child2] for p in event: assert p in all_particles all_particles.remove(p) assert all_particles == [] child3 = Particle(particle_id=Particle.Type.electron_antineutrino, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9) event.add_children(child1, child3) assert len(event) == 4 all_particles = [particle, child1, child2, child3] for p in event: assert p in all_particles all_particles.remove(p) assert all_particles == []
def test_total_cross_section_nubar(self, energy, sigma): """Test that total_cross_section attribute is correct for antineutrino interactions""" cc_nubar = Particle(particle_id=Particle.Type.electron_antineutrino, vertex=[100, 200, -500], direction=[0, 0, 1], energy=energy, interaction_model=CTWInteraction, interaction_type=Interaction.Type.charged_current) nc_nubar = Particle(particle_id=Particle.Type.electron_antineutrino, vertex=[100, 200, -500], direction=[0, 0, 1], energy=energy, interaction_model=CTWInteraction, interaction_type=Interaction.Type.neutral_current) undef_nubar = Particle(particle_id=Particle.Type.electron_antineutrino, vertex=[100, 200, -500], direction=[0, 0, 1], energy=energy, interaction_model=CTWInteraction, interaction_type=Interaction.Type.undefined) assert (cc_nubar.interaction.total_cross_section == nc_nubar.interaction.total_cross_section) assert (cc_nubar.interaction.cross_section + nc_nubar.interaction.cross_section == pytest.approx( cc_nubar.interaction.total_cross_section)) assert (undef_nubar.interaction.total_cross_section == cc_nubar.interaction.total_cross_section)
def test_get_exit_points(self, generator): """Test that the get_exit_points method returns accurate entry and exit""" p = Particle(particle_id=Particle.Type.electron_neutrino, vertex=(0, 0, -100), direction=(1, 0, 0), energy=1e9) enter_point, exit_point = generator.get_exit_points(p) assert np.array_equal(enter_point, [-2500, 0, -100]) assert np.array_equal(exit_point, [2500, 0, -100]) p = Particle(particle_id=Particle.Type.electron_neutrino, vertex=(0, 0, -100), direction=(0, 1, 0), energy=1e9) enter_point, exit_point = generator.get_exit_points(p) assert np.array_equal(enter_point, [0, -2500, -100]) assert np.array_equal(exit_point, [0, 2500, -100]) p = Particle(particle_id=Particle.Type.electron_neutrino, vertex=(0, 0, -100), direction=(0, 0, 1), energy=1e9) enter_point, exit_point = generator.get_exit_points(p) assert np.array_equal(enter_point, [0, 0, -3000]) assert np.array_equal(exit_point, [0, 0, 0]) p = Particle(particle_id=Particle.Type.electron_neutrino, vertex=(0, 0, -100), direction=(1, 1, 1), energy=1e9) enter_point, exit_point = generator.get_exit_points(p) assert np.array_equal(enter_point, [-2500, -2500, -2600]) assert np.array_equal(exit_point, [100, 100, 0])
def events(): """Fixture for forming basic list of Event objects""" return [ Event( Particle(particle_id=Particle.Type.electron_neutrino, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9)), Event( Particle(particle_id=Particle.Type.electron_antineutrino, vertex=[0, 0, 0], direction=[0, 0, -1], energy=1e9)), ]
def test_add_children_fails(self, particle): """Test that add_children fails if the parent isn't in the tree""" event = Event(particle) child1 = Particle(particle_id=Particle.Type.electron, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9, interaction_model=Interaction) child2 = Particle(particle_id=Particle.Type.positron, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9, interaction_model=Interaction) with pytest.raises(ValueError): event.add_children(child1, [child2])
def event(): """Fixture for forming basic Event object""" return Event( Particle(particle_id=Particle.Type.electron_neutrino, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9))
def _load_events(self): """ Pulls the next chunk of events into memory. Reads events up to the ``slice_range`` into memory from the current file. If the current file is exhausted, loads the next file. Returns ------- list List of `Event` objects read from the current file. Raises ------ StopIteration If the end of the last file in the file list has been reached. """ if self._file_index < 0 or self._event_index >= len(self._file): self._next_file() start = self._event_index stop = self._event_index + self.slice_range self._event_index += self.slice_range if stop > len(self._file): stop = len(self._file) self._events = [] self._event_counts = [] for file_event in self._file[start:stop]: info = file_event.get_particle_info() particles = [] for p in info: part = Particle(particle_id=p['particle_id'], vertex=(p['vertex_x'], p['vertex_y'], p['vertex_z']), direction=(p['direction_x'], p['direction_y'], p['direction_z']), energy=p['energy'], interaction_model=self.interaction_model, interaction_type=p['interaction_kind']) part.interaction.inelasticity = p['interaction_inelasticity'] part.interaction.em_frac = p['interaction_em_frac'] part.interaction.had_frac = p['interaction_had_frac'] part.survival_weight = p['survival_weight'] part.interaction_weight = p['interaction_weight'] particles.append(part) self._events.append(Event(particles)) self._event_counts.append(file_event.total_events_thrown)
def test_interaction_type_setting(self): """Test that the interaction type can be set during construction""" p = Particle(particle_id="nu_e", vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9, interaction_type=Interaction.Type.charged_current) assert p.interaction.kind == Interaction.Type.charged_current
def CTW_cc_nu(): """Fixture for forming a neutrino with charged-current CTWInteraction""" return Particle(particle_id=Particle.Type.electron_neutrino, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9, interaction_model=CTWInteraction, interaction_type=Interaction.Type.charged_current)
def CTW_nc_nubar(): """Fixture for forming a neutrino with neutral-current CTWInteraction""" return Particle(particle_id=Particle.Type.antielectron_neutrino, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9, interaction_model=CTWInteraction, interaction_type=Interaction.Type.neutral_current)
def test_get_parent(self, particle): """Test the ability to retrieve parent of a particle""" event = Event(particle) child1 = Particle(particle_id=Particle.Type.electron, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9, interaction_model=Interaction) child2 = Particle(particle_id=Particle.Type.positron, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9, interaction_model=Interaction) event.add_children(particle, [child1, child2]) assert event.get_parent(child1) == particle assert event.get_parent(child2) == particle assert event.get_parent(particle) is None
def test_get_from_level(self, particle): """Test the ability to get particles from a level in the tree""" event = Event(particle) child1 = Particle(particle_id=Particle.Type.electron, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9, interaction_model=Interaction) child2 = Particle(particle_id=Particle.Type.positron, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9, interaction_model=Interaction) event.add_children(particle, [child1, child2]) expected_level_0 = [particle] expected_level_1 = [child1, child2] for p in event.get_from_level(0): assert p in expected_level_0 expected_level_0.remove(p) assert expected_level_0 == [] for p in event.get_from_level(1): assert p in expected_level_1 expected_level_1.remove(p) assert expected_level_1 == [] assert event.get_from_level(2) == [] child3 = Particle(particle_id=Particle.Type.electron_antineutrino, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9) event.add_children(child1, child3) expected_level_0 = [particle] expected_level_1 = [child1, child2] expected_level_2 = [child3] for p in event.get_from_level(0): assert p in expected_level_0 expected_level_0.remove(p) assert expected_level_0 == [] for p in event.get_from_level(1): assert p in expected_level_1 expected_level_1.remove(p) assert expected_level_1 == [] for p in event.get_from_level(2): assert p in expected_level_2 expected_level_2.remove(p) assert expected_level_2 == []
def test_create_event(self, event): event2 = Event( Particle(particle_id="nu_e", vertex=[0, 0, 0], direction=[0, 0, -1], energy=1e9)) generator = ListGenerator([event, event2]) assert generator.create_event() == event assert generator.create_event() == event2
def test_direction_normalization(self): """Test that the particle direction is automatically normalized""" direction = [0, 1, 1] p = Particle(particle_id=Particle.Type.electron_neutrino, vertex=[100, 200, -500], direction=direction, energy=1e9) assert np.array_equal(p.direction, direction / np.linalg.norm(direction))
def test_get_exit_points(self, cyl_generator): """Test that the get_exit_points method returns appropriate exit points""" particle = Particle(particle_id='nu_e', energy=1e9, vertex=(0, 0, -1000), direction=(0, 0, 1)) with np.errstate(divide='ignore'): points = cyl_generator.get_exit_points(particle) assert np.array_equal(points[0], (0, 0, -3000)) assert np.array_equal(points[1], (0, 0, 0)) particle = Particle(particle_id='nu_e', energy=1e9, vertex=(0, 0, -1000), direction=(0, 0, -1)) with np.errstate(divide='ignore'): points = cyl_generator.get_exit_points(particle) assert np.array_equal(points[0], (0, 0, 0)) assert np.array_equal(points[1], (0, 0, -3000)) particle = Particle(particle_id='nu_e', energy=1e9, vertex=(0, 0, -1000), direction=(1, 0, 0)) points = cyl_generator.get_exit_points(particle) assert np.array_equal(points[0], (-5000, 0, -1000)) assert np.array_equal(points[1], (5000, 0, -1000)) particle = Particle(particle_id='nu_e', energy=1e9, vertex=(0, 0, -1000), direction=(0, 1, 0)) points = cyl_generator.get_exit_points(particle) assert np.array_equal(points[0], (0, -5000, -1000)) assert np.array_equal(points[1], (0, 5000, -1000)) particle = Particle(particle_id='nu_e', energy=1e9, vertex=(0, 0, -1000), direction=(1, 0, 1)) points = cyl_generator.get_exit_points(particle) assert np.array_equal(points[0], (-2000, 0, -3000)) assert np.array_equal(points[1], (1000, 0, 0)) particle = Particle(particle_id='nu_e', energy=1e9, vertex=(-4000, 0, -1000), direction=(1, 0, 1)) points = cyl_generator.get_exit_points(particle) assert np.array_equal(points[0], (-5000, 0, -2000)) assert np.array_equal(points[1], (-3000, 0, 0)) np.random.seed(SEED) edge2 = cyl_generator.dr**2 for _ in range(1000): particle = Particle(particle_id=cyl_generator.get_particle_type(), vertex=cyl_generator.get_vertex(), direction=cyl_generator.get_direction(), energy=cyl_generator.get_energy()) points = cyl_generator.get_exit_points(particle) for point in points: assert (point[0]**2 + point[1]**2 == pytest.approx(edge2) or point[2] == pytest.approx(0) or point[2] == pytest.approx(-cyl_generator.dz))
def test_get_exit_points(self, rect_generator): """Test that the get_exit_points method returns appropriate exit points""" particle = Particle(particle_id='nu_e', energy=1e9, vertex=(0, 0, -1000), direction=(0, 0, 1)) points = rect_generator.get_exit_points(particle) assert np.array_equal(points[0], (0, 0, -3000)) assert np.array_equal(points[1], (0, 0, 0)) particle = Particle(particle_id='nu_e', energy=1e9, vertex=(0, 0, -1000), direction=(0, 0, -1)) points = rect_generator.get_exit_points(particle) assert np.array_equal(points[0], (0, 0, 0)) assert np.array_equal(points[1], (0, 0, -3000)) particle = Particle(particle_id='nu_e', energy=1e9, vertex=(0, 0, -1000), direction=(1, 0, 0)) points = rect_generator.get_exit_points(particle) assert np.array_equal(points[0], (-2500, 0, -1000)) assert np.array_equal(points[1], (2500, 0, -1000)) particle = Particle(particle_id='nu_e', energy=1e9, vertex=(0, 0, -1000), direction=(0, 1, 0)) points = rect_generator.get_exit_points(particle) assert np.array_equal(points[0], (0, -2500, -1000)) assert np.array_equal(points[1], (0, 2500, -1000)) particle = Particle(particle_id='nu_e', energy=1e9, vertex=(0, 0, -1000), direction=(1, 0, 1)) points = rect_generator.get_exit_points(particle) assert np.array_equal(points[0], (-2000, 0, -3000)) assert np.array_equal(points[1], (1000, 0, 0)) particle = Particle(particle_id='nu_e', energy=1e9, vertex=(-2000, 0, -1000), direction=(1, 0, 1)) points = rect_generator.get_exit_points(particle) assert np.array_equal(points[0], (-2500, 0, -1500)) assert np.array_equal(points[1], (-1000, 0, 0)) np.random.seed(SEED) for _ in range(1000): particle = Particle(particle_id=rect_generator.get_particle_type(), vertex=rect_generator.get_vertex(), direction=rect_generator.get_direction(), energy=rect_generator.get_energy()) points = rect_generator.get_exit_points(particle) for point in points: assert (point[0] == pytest.approx(rect_generator.dx / 2) or point[0] == pytest.approx(-rect_generator.dx / 2) or point[1] == pytest.approx(rect_generator.dy / 2) or point[1] == pytest.approx(-rect_generator.dy / 2) or point[2] == pytest.approx(0) or point[2] == pytest.approx(-rect_generator.dz))
def kernel(): """Fixture for forming basic EventKernel object""" gen = ListGenerator( Event( Particle(particle_id="electron_neutrino", vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9))) return EventKernel(generator=gen, antennas=[Antenna(position=(0, 0, -100), noisy=False)])
def test_cross_section_nc_nubar(self, energy, sigma): """Test that cross_section method attribute is correct for neutral-current antineutrino interactions""" nc_nubar = Particle(particle_id=Particle.Type.electron_antineutrino, vertex=[100, 200, -500], direction=[0, 0, 1], energy=energy, interaction_model=CTWInteraction, interaction_type=Interaction.Type.neutral_current) assert nc_nubar.interaction.cross_section == pytest.approx(sigma, rel=0.001)
def test_cross_section_cc_nu(self, energy, sigma): """Test that cross_section attribute is correct for charged-current neutrino interactions""" cc_nu = Particle(particle_id=Particle.Type.electron_neutrino, vertex=[100, 200, -500], direction=[0, 0, 1], energy=energy, interaction_model=CTWInteraction, interaction_type=Interaction.Type.charged_current) assert cc_nu.interaction.cross_section == pytest.approx(sigma, rel=0.0001)
def test_get_children(self, particle): """Test the ability to retrieve children of a particle""" event = Event(particle) child1 = Particle(particle_id=Particle.Type.electron, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9, interaction_model=Interaction) child2 = Particle(particle_id=Particle.Type.positron, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9, interaction_model=Interaction) event.add_children(particle, [child1, child2]) expected_children = [child1, child2] for child in event.get_children(particle): assert child in expected_children expected_children.remove(child) assert expected_children == [] assert event.get_children(child1) == []
def test_get_weights(self, rect_generator): """Test that the get_weights returns floats between 0 and 1""" np.random.seed(SEED) for _ in range(1000): particle = Particle(particle_id=rect_generator.get_particle_type(), vertex=rect_generator.get_vertex(), direction=rect_generator.get_direction(), energy=rect_generator.get_energy()) weights = rect_generator.get_weights(particle) assert len(weights) == 2 assert 0 <= weights[0] <= 1 assert 0 <= weights[1] <= 1
def test_choose_interaction(self, energy): """Test that choose_interaction method properly samples the current ratio""" np.random.seed(SEED) cc_nu = Particle(particle_id=Particle.Type.electron_neutrino, vertex=[100, 200, -500], direction=[0, 0, 1], energy=energy, interaction_model=CTWInteraction) int_types = [] for _ in range(10000): int_types.append(cc_nu.interaction.choose_interaction().value - 1) fraction = 0.252162 + 0.0256 * np.log(np.log10(energy) - 1.76) assert np.mean(int_types) == pytest.approx(fraction, rel=0.05)
def test_choose_inelasticity_nc_nubar(self, energy, ybar): """Test that choose_inelasticity method properly samples the inelasticity distribution for neutral-current antineutrino interactions""" np.random.seed(SEED) nc_nubar = Particle(particle_id=Particle.Type.electron_antineutrino, vertex=[100, 200, -500], direction=[0, 0, 1], energy=energy, interaction_model=CTWInteraction, interaction_type=Interaction.Type.neutral_current) ys = [] for _ in range(10000): ys.append(nc_nubar.interaction.choose_inelasticity()) assert np.mean(ys) == pytest.approx(ybar, rel=0.2)
def test_id_coercion(self): """Test that the particle id can be set by enum value, int, or str""" p = Particle(particle_id=Particle.Type.electron, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9, interaction_model=Interaction) assert p.id == Particle.Type.electron p = Particle(particle_id=-12, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9) assert p.id == Particle.Type.electron_antineutrino p = Particle(particle_id="nu_mu", vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9) assert p.id == Particle.Type.muon_neutrino p = Particle(particle_id=None, vertex=[100, 200, -500], direction=[0, 0, 1], energy=1e9, interaction_model=Interaction) assert p.id == Particle.Type.undefined
def create_event(self): """ Generate a neutrino event in the ice volume. Creates a neutrino with a random vertex in the volume, a random direction, and an energy based on ``get_energy``. Particle type is randomly chosen, and its interaction type is also randomly chosen based on the branching ratio. Weights the particles according to their survival probability through the Earth and their probability of interacting in the ice at their vertex. If Earth shadowing has been turned on then particles which don't survive transit through the Earth are skipped, and surviving particles are given a survival weight of 1. Currently each `Event` returned consists of only a single `Particle`. Returns ------- Event Random neutrino event not shadowed by the Earth. See Also -------- pyrex.Event : Class for storing a tree of `Particle` objects representing an event. pyrex.Particle : Class for storing particle attributes. """ self.count += 1 vtx = self.get_vertex() u = self.get_direction() E = self.get_energy() particle_id = self.get_particle_type() particle = Particle(particle_id=particle_id, vertex=vtx, direction=u, energy=E, interaction_model=self.interaction_model) weights = self.get_weights(particle) if not self.shadow: particle.survival_weight = weights[0] particle.interaction_weight = weights[1] logger.debug( "Successfully created %s with survival weight %d and " + "interaction weight %d", particle, weights[0], weights[1]) return Event(particle) elif np.random.rand() < weights[0]: particle.survival_weight = 1 particle.interaction_weight = weights[1] logger.debug( "Successfully created %s with survival weight %d and " + "interaction weight %d", particle, weights[0], weights[1]) return Event(particle) else: # Particle was shadowed by the earth. Try again logger.debug("Particle creation shadowed by the Earth") return self.create_event()
def arz_pulse(): """Example Askaryan pulse from https://arxiv.org/pdf/1106.6283v3.pdf""" # Create particle to ensure shower energy is 3e9 GeV particle = Particle(particle_id=Particle.Type.electron_neutrino, vertex=(0, 0, -1000), direction=(0, 0, 1), energy=3e9, interaction_type="cc") particle.interaction.em_frac = 1 particle.interaction.had_frac = 0 n = ice.index(particle.vertex[2]) cherenkov_angle = np.arcsin(np.sqrt(1 - 1 / n**2)) return AskaryanSignal(times=np.linspace(0, 3e-9, 301), particle=particle, viewing_angle=cherenkov_angle - np.radians(0.3), t0=1e-9)
def create_event(self): """ Generate a neutrino. Pulls the next `Particle` object from the file(s) and places it into an `Event` by itself. Returns ------- Event Next neutrino `Event` object from the file(s). Raises ------ StopIteration If the end of the last file in the file list has been reached. See Also -------- pyrex.Event : Class for storing a tree of `Particle` objects representing an event. pyrex.Particle : Class for storing particle attributes. """ self._index += 1 if self.vertices is None or self._index >= len(self.vertices): self._next_file() return self.create_event() if self.interactions is None: interaction = None else: interaction = self.interactions[self._index] if self.weights is None: weight = 1 else: weight = self.weights[self._index] p = Particle(particle_id=self.ids[self._index], vertex=self.vertices[self._index], direction=self.directions[self._index], energy=self.energies[self._index], interaction_model=self.interaction_model, interaction_type=interaction, weight=weight) return Event(p)
def test_loop(self, event): """Test that the loop property allows for turning on and off the re-iteration of the list of events""" event2 = Event( Particle(particle_id="nu_e", vertex=[0, 0, 0], direction=[0, 0, -1], energy=1e9)) generator = ListGenerator([event, event2]) assert generator.create_event() == event assert generator.create_event() == event2 assert generator.create_event() == event assert generator.create_event() == event2 assert generator.create_event() == event generator = ListGenerator(event, loop=False) assert not generator.loop assert generator.create_event() == event with pytest.raises(StopIteration): generator.create_event()
def test_create_event(self, file_gen, tmpdir): """Test that create_event method loops over files correctly. Also tests ability to read files without explicit labels since test_particles_2.npz is created without explicit labels""" for i in range(4): event = file_gen.create_event() particle = event.roots[0] expected = Particle(particle_id=test_ids[i], vertex=test_vertices[i], direction=test_directions[i], energy=test_energies[i], interaction_type=test_interactions[i], weight=test_weights[i]) assert particle.id == expected.id assert np.array_equal(particle.vertex, expected.vertex) assert np.array_equal(particle.direction, expected.direction) assert particle.energy == expected.energy assert particle.interaction.kind == expected.interaction.kind assert particle.weight == expected.weight with pytest.raises(StopIteration): file_gen.create_event()
def test_weights_distribution(self, rect_generator): """Test that the get_weights distributions are as expected""" np.random.seed(SEED) survival_weights = [] interaction_weights = [] for _ in range(10000): particle = Particle(particle_id=rect_generator.get_particle_type(), vertex=rect_generator.get_vertex(), direction=rect_generator.get_direction(), energy=rect_generator.get_energy()) weights = rect_generator.get_weights(particle) survival_weights.append(weights[0]) interaction_weights.append(weights[1]) survival_hist, _ = np.histogram(survival_weights, range=(0, 1), bins=20) expected = np.array([ 4537.79, 74.1, 49.14, 43.61, 36.18, 29.1, 24.88, 20.71, 17.09, 14.46, 12.01, 9.93, 7.43, 4.97, 12.87, 27.77, 36.09, 47.87, 92.14, 4901.86 ]) assert np.allclose(survival_hist, expected, rtol=0.05, atol=10)