Merge pull request #279 from smash-transport/roch/fix_file_writer

Fix file writer Jetscape

src/sparkx/Jetscape.py

@@ -938,8 +938,9 @@ def get_sigmaGen(self):
 
     def print_particle_lists_to_file(self, output_file):
         """
-        Prints the current Jetscape data to an output file specified by :code:`output_file`
-        with the same format as the input file
+        Prints the current Jetscape data to an output file specified by
+        :code:`output_file` with the same format as the input file.
+        For empty events, only the event header is printed.
 
         Parameters
         ----------
@@ -966,9 +967,10 @@ def print_particle_lists_to_file(self, output_file):
                     f_out.write(header)
 
                     # Write the particle data to the file
-                    np.savetxt(
-                        f_out, particle_output, fmt="%d %d %d %g %g %g %g"
-                    )
+                    if particle_output.shape[0] != 0:
+                        np.savetxt(
+                            f_out, particle_output, fmt="%d %d %d %g %g %g %g"
+                        )
             else:
                 event = 0
                 num_out = self.num_output_per_event_[0][1]
@@ -979,7 +981,10 @@ def print_particle_lists_to_file(self, output_file):
                 f_out.write(header)
 
                 # Write the particle data to the file
-                np.savetxt(f_out, particle_output, fmt="%d %d %d %g %g %g %g")
+                if particle_output.shape[0] != 0:
+                    np.savetxt(
+                        f_out, particle_output, fmt="%d %d %d %g %g %g %g"
+                    )
 
             # Write the last line
             last_line = self.__get_last_line(self.PATH_JETSCAPE_) + "\n"

src/sparkx/Oscar.py

@@ -1123,8 +1123,9 @@ def multiplicity_cut(self, min_multiplicity):
 
     def print_particle_lists_to_file(self, output_file):
         """
-        Prints the current Oscar data to an output file specified by :code:`output_file`
-        with the same format as the input file
+        Prints the current Oscar data to an output file specified by
+        :code:`output_file` with the same format as the input file.
+        For empty events, only the event header and footer are printed.
 
         Parameters
         ----------

tests/test_Jetscape.py

@@ -41,6 +41,14 @@ def jetscape_file_path_partons():
     )
 
 
+@pytest.fixture
+def jetscape_file_no_hadrons():
+    # Assuming your test file is in the same directory as test_files/
+    return os.path.join(
+        os.path.dirname(__file__), "test_files", "test_jetscape_no_hadrons.dat"
+    )
+
+
 def create_temporary_jetscape_file(
     path, num_events, output_per_event_list=None
 ):
@@ -457,6 +465,19 @@ def test_Jetscape_print(jetscape_file_path, output_path):
     os.remove(output_path)
 
 
+def test_Jetscape_print_with_empty_events(
+    jetscape_file_path, output_path, jetscape_file_no_hadrons
+):
+    jetscape = Jetscape(jetscape_file_path)
+    jetscape.particle_list_ = [[], [], [], [], []]
+    jetscape.num_output_per_event_ = np.array(
+        [[1, 0], [2, 0], [3, 0], [4, 0], [5, 0]]
+    )
+    jetscape.print_particle_lists_to_file(output_path)
+    assert filecmp.cmp(jetscape_file_no_hadrons, output_path)
+    os.remove(output_path)
+
+
 def test_Jetscape_get_sigmaGen(jetscape_file_path):
     jetscape = Jetscape(jetscape_file_path)
     assert jetscape.get_sigmaGen() == (0.000314633, 6.06164e-07)

tests/test_files/test_jetscape_no_hadrons.dat

@@ -0,0 +1,7 @@
+#	JETSCAPE_FINAL_STATE	v2	|	N	pid	status	E	Px	Py	Pz
+#	Event	1	weight	1	EPangle	0	N_hadrons	0
+#	Event	2	weight	1	EPangle	0	N_hadrons	0
+#	Event	3	weight	1	EPangle	0	N_hadrons	0
+#	Event	4	weight	1	EPangle	0	N_hadrons	0
+#	Event	5	weight	1	EPangle	0	N_hadrons	0
+#	sigmaGen	0.000314633	sigmaErr	6.06164e-07