Skip to content

Commit

Permalink
Add tests and fix event counting for Jetscape
Browse files Browse the repository at this point in the history
Fix num_output_per_event() so that it starts counting from 0 for Oscar
and from 1 for Jetscape files. Also add test for warning in case of
empty event list.
  • Loading branch information
Lucas Constantin committed Nov 7, 2024
1 parent 504bc7b commit 799fa39
Show file tree
Hide file tree
Showing 5 changed files with 19 additions and 7 deletions.
2 changes: 1 addition & 1 deletion src/sparkx/BaseStorer.py
Original file line number Diff line number Diff line change
Expand Up @@ -560,7 +560,7 @@ def _update_num_output_per_event_after_filter(self) -> None:
# Handle the case where num_output_per_event_ is a two-dimensional array
updated_num_output_per_event = np.ndarray((len(self.particle_list_),2), dtype=int)
for event in range(len(self.particle_list_)):
updated_num_output_per_event[event][0] = event
updated_num_output_per_event[event][0] = event + self.num_output_per_event_[0][0]
updated_num_output_per_event[event][1] = len(
self.particle_list_[event]
)
Expand Down
4 changes: 2 additions & 2 deletions src/sparkx/Filter.py
Original file line number Diff line number Diff line change
Expand Up @@ -797,15 +797,15 @@ def spacetime_rapidity_cut(particle_list, cut_value):

def multiplicity_cut(particle_list, cut_value_tuple):
"""
Apply multiplicity cut. Remove all events with a multiplicity not coplying
Apply multiplicity cut. Remove all events with a multiplicity not complying
with cut_value_tuple.
Parameters
----------
particle_list:
List with lists containing particle objects for the events
multiplicity_range : tuple
cut_value_tuple : tuple
Upper and lower bound for multiplicity. If the multiplicity of an event is
not in this range, the event is discarded. The range is inclusive on the
lower bound and exclusive on the upper bound.
Expand Down
6 changes: 4 additions & 2 deletions src/sparkx/Jetscape.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ class Jetscape(BaseStorer):
>>> jetscape = Jetscape(JETSCAPE_FILE_PATH)
>>>
>>> pions = jetscape.multiplicity_cut(500).participants().particle_species((211, -211, 111))
>>> pions = jetscape.multiplicity_cut(500, None).participants().particle_species((211, -211, 111))
>>>
>>> # save the pions of all events as nested list
>>> pions_list = pions.particle_list()
Expand All @@ -150,7 +150,7 @@ class Jetscape(BaseStorer):
Let's assume we only want to keep pions in events with a
multiplicity > 500:
>>> jetscape = Jetscape(JETSCAPE_FILE_PATH, filters={'multiplicity_cut':500, 'particle_species':(211, -211, 111)}})
>>> jetscape = Jetscape(JETSCAPE_FILE_PATH, filters={'multiplicity_cut':(500,None), 'particle_species':(211, -211, 111)}})
>>>
>>> # print the pions to a jetscape file
>>> jetscape.print_particle_lists_to_file('./particle_lists.dat')
Expand Down Expand Up @@ -304,6 +304,8 @@ def print_particle_lists_to_file(self, output_file: str) -> None:
raise ValueError("The number of output per event is empty.")
if self.num_events_ is None:
raise ValueError("The number of events is empty.")
if self.num_events_ == 0:
raise ValueError("The number of events is zero.")

# Open the output file with buffered writing (25 MB)
with open(output_file, "w", buffering=25 * 1024 * 1024) as f_out:
Expand Down
4 changes: 2 additions & 2 deletions src/sparkx/Oscar.py
Original file line number Diff line number Diff line change
Expand Up @@ -134,7 +134,7 @@ class Oscar(BaseStorer):
>>> oscar = Oscar(OSCAR_FILE_PATH)
>>>
>>> pions = oscar.multiplicity_cut(500).participants().particle_species((211, -211, 111))
>>> pions = oscar.multiplicity_cut(500, None).participants().particle_species((211, -211, 111))
>>>
>>> # save the pions of all events as nested list
>>> pions_list = pions.particle_list()
Expand All @@ -154,7 +154,7 @@ class Oscar(BaseStorer):
Let's assume we only want to keep participant pions in events with a
multiplicity > 500:
>>> oscar = Oscar(OSCAR_FILE_PATH, filters={'multiplicity_cut':500, 'participants':True, 'particle_species':(211, -211, 111)})
>>> oscar = Oscar(OSCAR_FILE_PATH, filters={'multiplicity_cut':(500,None), 'participants':True, 'particle_species':(211, -211, 111)})
>>>
>>> # print the pions to an oscar file
>>> oscar.print_particle_lists_to_file('./particle_lists.oscar')
Expand Down
10 changes: 10 additions & 0 deletions tests/test_Oscar.py
Original file line number Diff line number Diff line change
Expand Up @@ -684,6 +684,16 @@ def test_standard_oscar_print(tmp_path, output_path):
os.remove(output_path)


def test_empty_oscar_print(tmp_path, output_path):
tmp_oscar_file = create_temporary_oscar_file(
tmp_path, 5, "Oscar2013", [1, 7, 0, 36, 5]
)
oscar = Oscar(tmp_oscar_file).multiplicity_cut((100000000,None))
with pytest.warns(UserWarning):
oscar.print_particle_lists_to_file(output_path)
os.remove(output_path)


def test_extended_oscar_print(tmp_path, output_path):
tmp_oscar_file = create_temporary_oscar_file(
tmp_path, 5, "Oscar2013Extended", [4, 1, 42, 0, 3]
Expand Down

0 comments on commit 799fa39

Please sign in to comment.