Merge pull request #181 from smash-transport/roch/add_JS_partons

Add JS partons

CHANGELOG.md

@@ -23,6 +23,8 @@ Date: XXX
 
 ### Added
 * ReactionPlaneFlow: Returns also the angle, not only the flow value
+* Jetscape: Possibility to read in parton output files with optional `particletype` parameter in constructor
+* Particle: Include changes for parton read in from Jetscape class, multiply quark charges by 3 to make them integers
 
 
 [Link to diff from previous version](https://github.com/smash-transport/sparkx/compare/v1.2.1...v1.2.2)

src/sparkx/Jetscape.py

@@ -60,6 +60,11 @@ class Jetscape:
                 memory. The names of the filters are the same as the names of |br|
                 the filter methods. All filters are applied in the order in |br|
                 which they appear in the dictionary.
+            * - :code:`particletype` (str)
+              - This parameter allows to switch from the standard hadron file |br|
+                to the parton output of JETSCAPE. The parameter can be set to |br|
+                :code:`particletype='hadron'` (default) or :code:`particletype='parton'`.
+                Quark charges are multiplied by 3 to make them integer values.
 
         .. |br| raw:: html
 
@@ -191,10 +196,11 @@ def __init__(self, JETSCAPE_FILE, **kwargs):
         self.num_output_per_event_ = None
         self.num_events_ = None
         self.particle_list_ = None
+        self.particle_type_ = 'hadron'
         self.optional_arguments_ = kwargs
 
         for keys in self.optional_arguments_.keys():
-            if keys not in ['events', 'filters']:
+            if keys not in ['events', 'filters', 'particletype']:
                 raise ValueError('Unknown keyword argument used in constructor')
 
         if 'events' in self.optional_arguments_.keys() and isinstance(self.optional_arguments_['events'], tuple):
@@ -206,6 +212,19 @@ def __init__(self, JETSCAPE_FILE, **kwargs):
         elif 'events' in self.optional_arguments_.keys() and isinstance(self.optional_arguments_['events'], int):
             if self.optional_arguments_['events'] < 0:
                 raise ValueError('Event number must be positive')
+
+        if 'particletype' in self.optional_arguments_.keys() and isinstance(self.optional_arguments_['particletype'], str):
+            if (self.optional_arguments_['particletype'] == 'hadron') or (self.optional_arguments_['particletype'] == 'parton'):
+                self.particle_type_ = self.optional_arguments_['particletype']
+            else:
+                raise ValueError("'particletype' has to be 'hadron' or 'parton'")
+        elif 'particletype' in self.optional_arguments_.keys() and not isinstance(self.optional_arguments_['particletype'], str):
+            raise TypeError("'particletype' is not given as a string value") 
+
+        if self.particle_type_ == 'hadron':
+            self.particle_type_defining_string_ = 'N_hadrons'
+        else:
+            self.particle_type_defining_string_ = 'N_partons'
 
         self.set_num_output_per_event()
         self.set_particle_list(kwargs)
@@ -432,7 +451,7 @@ def set_num_output_per_event(self):
                 line = jetscape_file.readline()
                 if not line:
                     break
-                elif '#' in line and 'N_hadrons' in line:
+                elif '#' in line and self.particle_type_defining_string_ in line:
                     line_str = line.replace('\n','').replace('\t',' ').split(' ')
                     event = line_str[2]
                     num_output = line_str[8]
@@ -849,7 +868,7 @@ def print_particle_lists_to_file(self, output_file):
                         data_to_write.append(header_file)
 
                     # Header for each event
-                    header = f'#\tEvent\t{event}\tweight\t1\tEPangle\t0\tN_hadrons\t{num_out}\n'
+                    header = f'#\tEvent\t{event}\tweight\t1\tEPangle\t0\t{self.particle_type_defining_string_}\t{num_out}\n'
                     data_to_write.append(header)
 
                     # Convert particle data to formatted strings
@@ -867,7 +886,7 @@ def print_particle_lists_to_file(self, output_file):
                     data_to_write.append(header_file)
 
                 # Header for each event
-                header = f'#\tEvent\t{event}\tweight\t1\tEPangle\t0\tN_hadrons\t{num_out}\n'
+                header = f'#\tEvent\t{event}\tweight\t1\tEPangle\t0\t{self.particle_type_defining_string_}\t{num_out}\n'
                 data_to_write.append(header)
 
                 # Convert particle data to formatted strings

src/sparkx/Particle.py

@@ -200,7 +200,7 @@ class Particle:
         :linenos:
 
         >>> particle_quantity_JETSCAPE = np.array([0,2114,11,2.01351754,1.30688601,-0.422958786,-0.512249773])
-        >>> particle = Particle(input_format="JETSCAPE", particle_array=particle_array_oscar2013)
+        >>> particle = Particle(input_format="JETSCAPE", particle_array=particle_quantity_JETSCAPE)
 
     Notes
     -----
@@ -209,6 +209,11 @@ class Particle:
     All quantities are saved in a numpy array member variable `data_`. The datatype
     of this array is float, therefore casting is required when int or bool values are 
     required.
+
+    When JETSCAPE creates particle objects, which are partons, the charge is multiplied
+    by 3 to make it an integer.
+    The functions `is_strange()` and `is_heavy_flavor()` should not be used in this
+    case.
     """
     __slots__ = ['data_'] 
     def __init__(self,input_format=None,particle_array=None):
@@ -252,6 +257,8 @@ def __initialize_from_array(self,input_format,particle_array):
 
         If the input format is "JETSCAPE," additional attributes (mass_ and 
         charge_) are computed based on energy-momentum and PDG code.
+        If the JETSCAPE class is created with parton output, then the charge is
+        multiplied by 3 to make it an integer.
 
         """
         #first entry: index in data array
@@ -566,6 +573,10 @@ def charge(self):
 
     @charge.setter
     def charge(self,value):
+        # this is for the case a parton is created from the JETSCAPE reader
+        # handle quarks with 3 times the charge to make it integer
+        if np.abs(value) < 1:
+            value *= 3
         self.data_[12] = value
 
     @property

tests/test_Jetscape.py

@@ -26,6 +26,11 @@ def jetscape_file_path():
     # Assuming your test file is in the same directory as test_files/
     return os.path.join(os.path.dirname(__file__), 'test_files', 'test_jetscape.dat')
 
+@pytest.fixture
+def jetscape_file_path_partons():
+    # Assuming your test file is in the same directory as test_files/
+    return os.path.join(os.path.dirname(__file__), 'test_files', 'test_jetscape_partons.dat')
+
 def create_temporary_jetscape_file(path, num_events, output_per_event_list=None):
     """
     This function creates a sample Jetscape file "jetscape_test.dat" in the 
@@ -383,3 +388,20 @@ def test_Jetscape_charge_filter_one_event(jetscape_file_path):
     jetscape = Jetscape(jetscape_file_path, events=0).charged_particles()
     assert jetscape.num_events() == 1
     assert (jetscape.num_output_per_event() == np.array([[1, 14]])).all()
+
+def test_Jetscape_read_parton_file(jetscape_file_path_partons):
+    jetscape = Jetscape(jetscape_file_path_partons,particletype='parton')
+    assert jetscape.particle_type_ == 'parton'
+    assert jetscape.num_events() == 5
+
+    charged_quarks = jetscape.charged_particles().particle_objects_list()
+    for ev in range(5):
+        assert len(charged_quarks[ev]) == 2
+
+    # Test construction with wrong particletype
+    with pytest.raises(ValueError):
+        Jetscape(jetscape_file_path_partons,particletype='wrong')
+
+    # Test construction with wrong particletype (not string)
+    with pytest.raises(TypeError):
+        Jetscape(jetscape_file_path_partons,particletype=1)

tests/test_Particle.py

@@ -798,4 +798,17 @@ def test_spin_degeneracy_from_pdg_invalid_values():
 
     result = p.spin_degeneracy()
 
-    assert np.isnan(result)
+    assert np.isnan(result)
+
+def test_create_quark_from_pdg_valid_values():
+    particle_quantity_JETSCAPE = np.array([0,1,11,2.01351754,1.30688601,-0.422958786,-0.512249773])
+    particle = Particle(input_format="JETSCAPE", particle_array=particle_quantity_JETSCAPE)
+
+    assert particle.charge == -1
+    assert math.isclose(particle.mass, 1.38, rel_tol=1e-3)
+
+def test_create_gluon_from_pdg_valid_values():
+    particle_quantity_JETSCAPE = np.array([0,21,11,2.01351754,1.30688601,-0.422958786,-0.512249773])
+    particle = Particle(input_format="JETSCAPE", particle_array=particle_quantity_JETSCAPE)
+
+    assert particle.charge == 0

tests/test_files/test_jetscape_partons.dat

@@ -0,0 +1,40 @@
+#	JETSCAPE_FINAL_STATE	v2	|	N	pid	status	E	Px	Py	Pz
+#	Event	1	weight	1	EPangle	0	N_partons	6
+0 5 0 37.8629 -24.9369 -22.021 -17.2927
+1 21 0 3.3561 -1.59367 -2.42797 -1.65162
+2 21 0 4.97057 2.06545 -1.47451 -4.17807
+3 -5 0 32.5379 15.7891 23.3729 15.3411
+4 21 0 1.27252 1.2153 -0.226852 -0.0450638
+5 21 0 11.2 7.46077 2.77751 7.82638
+#	Event	2	weight	1	EPangle	0	N_partons	7
+0 21 0 6.5714 -3.58228 5.21497 0.709817
+1 1 0 35.9744 -29.4141 18.2972 -9.422
+2 21 0 1.40696 -0.510795 1.03982 0.766398
+3 -1 0 28.6536 26.1888 -8.98468 7.04292
+4 21 0 1.31288 -0.719053 -0.927709 -0.534268
+5 21 0 1.80355 1.20636 -0.789226 0.990258
+6 21 0 15.476 6.83107 -13.8504 0.446872
+#	Event	3	weight	1	EPangle	0	N_partons	5
+0 1 0 44.2829 1.19823 19.716 -39.6233
+1 21 0 2.98482 1.32127 1.75765 1.96571
+2 21 0 7.19202 -1.66465 -3.86631 5.76157
+3 -1 0 34.5746 -0.440891 -17.3786 29.8002
+4 21 0 2.16449 -0.413959 -0.228692 2.09588
+#	Event	4	weight	1	EPangle	0	N_partons	8
+0 21 0 2.20421 -1.52689 0.320971 -1.5458
+1 21 0 4.21275 3.57545 1.76868 -1.0125
+2 21 0 0.8271 0.0189902 0.246713 0.781088
+3 21 0 2.50326 2.23249 1.01671 -0.384335
+4 3 0 17.4496 14.7309 7.49091 5.52849
+5 21 0 20.5896 17.3009 10.2522 4.32368
+6 -3 0 42.4491 -36.5628 -20.3065 -7.2059
+7 21 0 0.961492 0.230912 -0.789708 -0.48473
+#	Event	5	weight	1	EPangle	0	N_partons	7
+0 21 0 3.52122 -0.873291 0.80143 -3.30805
+1 21 0 1.24819 -0.696865 0.873414 -0.540979
+2 1 0 26.2539 -8.4428 -5.10896 24.2845
+3 21 0 17.9126 -6.01729 -1.60042 16.7715
+4 21 0 1.09568 -0.545857 0.938975 0.077616
+5 -1 0 35.3418 14.1926 2.67511 -32.188
+6 21 0 5.82524 2.38357 1.42046 -5.09658
+#	sigmaGen	0.000285735	sigmaErr	2.47804e-05