diff --git a/.github/workflows/code_formatting.yml b/.github/workflows/code_formatting.yml index cba24b31..22777bbb 100644 --- a/.github/workflows/code_formatting.yml +++ b/.github/workflows/code_formatting.yml @@ -5,22 +5,20 @@ on: branches: - main - sparkx_devel - types: - - closed # Trigger when the PR is closed (includes merging) jobs: code-formatting: - if: github.event.pull_request.merged == true runs-on: ubuntu-latest steps: - - name: Checkout repository - uses: actions/checkout@v2 + - name: Checkout PR branch + uses: actions/checkout@v3 with: - ref: ${{ github.sha }} # Check out the exact commit of the merged PR + ref: ${{ github.head_ref }} + fetch-depth: 0 # Ensures full history is available - name: Set up Python - uses: actions/setup-python@v2 + uses: actions/setup-python@v4 with: python-version: '3.9' @@ -29,22 +27,25 @@ jobs: python -m pip install --upgrade pip pip install black==24.8.0 - - name: Format code + - name: Run black to check formatting + id: black-check run: | - black --line-length 80 src/sparkx tests/ + black --check --line-length 80 src/sparkx tests/ + continue-on-error: true - - name: Commit changes - run: | - git config --local user.name "GitHub Action" - git config --local user.email "action@github.com" - git add src/sparkx tests/ - git commit -m "Automatically format code using black" || echo "No changes to commit" + - name: Capture formatting status + if: ${{ steps.black-check.outcome == 'failure' }} + run: echo "needs_formatting=true" >> $GITHUB_ENV - - name: Push changes - run: | - git push origin ${{ github.head_ref }} + - name: Format code with black + if: env.needs_formatting == 'true' + run: black --line-length 80 src/sparkx tests/ - - name: Run format test again + - name: Push formatted changes + if: env.needs_formatting == 'true' run: | - # Add your test commands here - black --check --line-length 80 src/sparkx tests/ + git config --global user.name "github-actions[bot]" + git config --global user.email "github-actions[bot]@users.noreply.github.com" + git add src/sparkx tests/ + git commit -m "Auto-format code with black" + git push origin ${{ github.head_ref }} diff --git a/CHANGELOG.md b/CHANGELOG.md index c8c5e4f9..dc9696c6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,7 @@ Date: * Added support for the SMASH 3.2 feature of custom output format * Add option to add two Oscar/Jetscape/ParticleObjectStorer instances while preserving the event order * BulkObservables: Add a class for calculating spectra and integrated yields +* Oscar: Add function to extract the impact parameters ### Changed * Particle: Rename several methods for a more intuitive naming scheme. Renamed methods are: diff --git a/pyproject.toml b/pyproject.toml index b82a04c8..eedd4e64 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,9 +6,11 @@ build-backend = "setuptools.build_meta" name = "sparkx" version = "2.0.0" authors = [ + { name="Lucas Constantin", email="constantin@itp.uni-frankfurt.de"}, { name="Niklas Götz", email="goetz@itp.uni-frankfurt.de"}, { name="Renata Krupczak", email="rkrupczak@physik.uni-bielefeld.de"}, { name="Hendrik Roch", email="hroch@wayne.edu" }, + { name="Carl Rosenkvist", email="rosenkvist@itp.uni-frankfurt.de"}, { name="Nils Sass", email="nsass@itp.uni-frankfurt.de" } ] description = "Software Package for Analyzing Relativistic Kinematics in Collision eXperiments" @@ -25,11 +27,10 @@ classifiers = [ ] dependencies = [ - "particle==0.23.0", + "particle>=0.23.0", "numpy>=1.23.5", "scipy>=1.10.1", - "abc-property==1.0", - "fastjet==3.4.1.3", + "fastjet>=3.4.2.1", "matplotlib>=3.7.1" ] diff --git a/requirements.txt b/requirements.txt index 48a4ca2b..4eca8390 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,8 @@ # Main dependencies -particle==0.23.0 +particle>=0.23.0 numpy>=1.23.5 scipy>=1.10.1 -abc-property==1.0 -fastjet==3.4.1.3 +fastjet>=3.4.2.1 matplotlib>=3.7.1 # Documentation dependencies @@ -15,7 +14,7 @@ setuptools>=68.0.0 # Development and testing dependencies pytest>=7.4.3 -black>=24.8.0 +black==24.8.0 # Editable Sparkx install -e . diff --git a/setup.py b/setup.py index 721f303c..ae1d926e 100644 --- a/setup.py +++ b/setup.py @@ -5,18 +5,17 @@ packages=find_packages(where='src'), package_dir={'': 'src'}, install_requires=[ - "particle==0.23.0", + "particle>=0.23.0", "numpy>=1.23.5", "scipy>=1.10.1", - "abc-property==1.0", - "fastjet==3.4.1.3", + "fastjet>=3.4.2.1", "matplotlib>=3.7.1", ], version='2.0.0', description='Software Package for Analyzing Relativistic Kinematics in Collision eXperiments', long_description=open("README.md").read(), long_description_content_type="text/markdown", - author='Niklas Götz, Renata Krupczak, Hendrik Roch, Nils Sass', + author='Lucas Constantin, Niklas Götz, Renata Krupczak, Hendrik Roch, Carl Rosenkvist, Nils Sass', author_email="goetz@itp.uni-frankfurt.de, hroch@wayne.edu, nsass@itp.uni-frankfurt.de", url="https://smash-transport.github.io/sparkx/", download_url="https://github.com/smash-transport/sparkx", diff --git a/src/sparkx/BaseStorer.py b/src/sparkx/BaseStorer.py index 599a7aef..5571ddfc 100644 --- a/src/sparkx/BaseStorer.py +++ b/src/sparkx/BaseStorer.py @@ -23,7 +23,7 @@ class BaseStorer(ABC): ---------- num_output_per_event_ : numpy.array Array containing the event number and the number of particles in this - event as :code:`num_output_per_event_[event i][num_output in event i]` + event as :code:`num_output_per_event_[event i][num_output in event i]` (updated when filters are applied) num_events_ : int Number of events contained in the Oscar object (updated when filters @@ -105,11 +105,11 @@ def __init__( self.particle_list_, self.num_events_, self.num_output_per_event_, - self.custom_attr_list + self.custom_attr_list, ) = self.loader_.load(**kwargs) else: raise ValueError("Loader has not been created properly") - + def __add__(self, other: "BaseStorer") -> "BaseStorer": """ Adds two BaseStorer objects by combining their particle lists and updating num_output_per_event accordingly. @@ -134,12 +134,14 @@ def __add__(self, other: "BaseStorer") -> "BaseStorer": """ if not isinstance(other, BaseStorer): raise TypeError("Can only add BaseStorer objects") - + # Ensure that both instances are of the same class if type(self) is not type(other): raise TypeError("Can only add objects of the same class") - combined_particle_list: list = self.particle_list_ + other.particle_list_ + combined_particle_list: list = ( + self.particle_list_ + other.particle_list_ + ) # Ensure num_output_per_event_ is not None if self.num_output_per_event_ is None: @@ -156,18 +158,22 @@ def __add__(self, other: "BaseStorer") -> "BaseStorer": ) # Adjust event_number for the parts that originally belonged to other - combined_num_output_per_event[self.num_events_:, 0] += self.num_events_ + combined_num_output_per_event[self.num_events_ :, 0] += self.num_events_ combined_storer: BaseStorer = self.__class__.__new__(self.__class__) - combined_storer.__dict__.update(self.__dict__) # Inherit all properties from self + combined_storer.__dict__.update( + self.__dict__ + ) # Inherit all properties from self combined_storer._update_after_merge(other) combined_storer.particle_list_ = combined_particle_list combined_storer.num_output_per_event_ = combined_num_output_per_event combined_storer.num_events_ = self.num_events_ + other.num_events_ - combined_storer.loader_ = None # Loader is not applicable for combined object + combined_storer.loader_ = ( + None # Loader is not applicable for combined object + ) return combined_storer - + @abstractmethod def _update_after_merge(self, other: "BaseStorer") -> None: """ @@ -203,7 +209,7 @@ def num_output_per_event(self) -> Optional[np.ndarray]: :code:`num_output_per_event[event_n, number_of_particles_in_event_n]` - :code:`num_output_per_event` is updated with every manipulation e.g. + :code:`num_output_per_event` is updated with every manipulation e.g. after applying cuts. Returns @@ -336,7 +342,7 @@ def particle_species( Returns ------- self : BaseStorer object - Containing only particle species specified by :code:`pdg_list` for + Containing only particle species specified by :code:`pdg_list` for every event """ self.particle_list_ = particle_species(self.particle_list_, pdg_list) @@ -348,7 +354,7 @@ def remove_particle_species( self, pdg_list: Union[int, Union[Tuple[int], List[int], np.ndarray]] ) -> "BaseStorer": """ - Remove particle species from :code:`particle_list` by their PDG ID in + Remove particle species from :code:`particle_list` by their PDG ID in every event. Parameters @@ -409,7 +415,7 @@ def lower_event_energy_cut( Parameters ---------- minimum_event_energy : int or float - The minimum event energy threshold. Should be a positive integer or + The minimum event energy threshold. Should be a positive integer or float. Returns @@ -421,10 +427,10 @@ def lower_event_energy_cut( Raises ------ TypeError - If the :code:`minimum_event_energy` parameter is not an integer or + If the :code:`minimum_event_energy` parameter is not an integer or float. ValueError - If the :code:`minimum_event_energy` parameter is less than or + If the :code:`minimum_event_energy` parameter is less than or equal to 0. """ self.particle_list_ = lower_event_energy_cut( @@ -502,7 +508,7 @@ def rapidity_cut( cut_value : float If a single value is passed, the cut is applied symmetrically around 0. - For example, if :code:`cut_value = 1`, only particles with rapidity + For example, if :code:`cut_value = 1`, only particles with rapidity in :code:`[-1.0, 1.0]` are kept. cut_value : tuple @@ -556,7 +562,7 @@ def spacetime_rapidity_cut( ) -> "BaseStorer": """ Apply spacetime rapidity (space-time rapidity) cut to all events and - remove all particles with spacetime rapidity not complying with + remove all particles with spacetime rapidity not complying with cut_value. Parameters @@ -587,7 +593,7 @@ def spacetime_rapidity_cut( def multiplicity_cut( self, cut_value_tuple: Tuple[Union[float, None], Union[float, None]] - ) -> "BaseStorer": + ) -> "BaseStorer": """ Apply multiplicity cut. Remove all events with a multiplicity not complying with cut_value. @@ -596,7 +602,7 @@ def multiplicity_cut( ---------- 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 + not in this range, the event is discarded. The range is inclusive on the lower bound and exclusive on the upper bound. Returns @@ -634,7 +640,7 @@ def spacetime_cut( Returns ------- self : BaseStorer object - Containing only particles complying with the spacetime cut for all + Containing only particles complying with the spacetime cut for all events """ self.particle_list_ = spacetime_cut( @@ -662,7 +668,7 @@ def particle_status( Returns ------- self : BaseStorer object - Containing only hadrons with status specified by + Containing only hadrons with status specified by :code:`status_list` for every event """ self.particle_list_ = particle_status(self.particle_list_, status_list) @@ -848,9 +854,13 @@ def _update_num_output_per_event_after_filter(self) -> None: self.num_output_per_event_[1] = len(self.particle_list_[0]) elif self.num_output_per_event_.ndim == 2: # Handle the case where num_output_per_event_ is a two-dimensional array - updated_num_output_per_event : np.ndarray = np.ndarray((len(self.particle_list_),2), dtype=int) + updated_num_output_per_event: np.ndarray = 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 + self.num_output_per_event_[0][0] + 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] ) @@ -869,7 +879,7 @@ def print_particle_lists_to_file(self, output_file: str) -> None: Prints the particle lists to a specified file. This method should be implemented by subclasses to print the particle - lists to the specified output file. The method raises a + lists to the specified output file. The method raises a :code:`NotImplementedError` if it is not overridden by a subclass. Parameters diff --git a/src/sparkx/CentralityClasses.py b/src/sparkx/CentralityClasses.py index 25928d08..aa7cecd6 100644 --- a/src/sparkx/CentralityClasses.py +++ b/src/sparkx/CentralityClasses.py @@ -16,6 +16,12 @@ class CentralityClasses: """ Class for defining centrality classes based on event multiplicity. + .. note:: + It is the user's responsibility to ensure that the amount of events used + to determine the centrality classes is sufficient to provide reliable + results. The recommended minimum number of events is at least a few + hundred. + Parameters ---------- events_multiplicity : list or numpy.ndarray @@ -27,7 +33,7 @@ class CentralityClasses: Raises ------ TypeError - If :code:`events_multiplicity` or :code:`centrality_bins` is not a list + If :code:`events_multiplicity` or :code:`centrality_bins` is not a list or :code:`numpy.ndarray`. Attributes @@ -58,10 +64,17 @@ class CentralityClasses: .. code-block:: python :linenos: - >>> centrality_obj = CentralityClasses(events_multiplicity=[10, 15, 20, 25], - ... centrality_bins=[0, 25, 50, 75, 100]) - >>> centrality_obj.get_centrality_class(18) - 1 + >>> from sparkx import * + >>> import random as rd + >>> rd.seed(0) + + >>> # generate 300 random event multiplicities between 50 and 1500 + >>> events_multiplicity = [rd.randint(50, 1500) for _ in range(300)] + >>> centrality_bins = [0, 10, 30, 50, 70, 90, 100] + >>> centrality_obj = CentralityClasses(events_multiplicity=events_multiplicity, + ... centrality_bins=centrality_bins) + >>> centrality_obj.get_centrality_class(1490) + 0 >>> centrality_obj.output_centrality_classes('centrality_output.txt') """ @@ -220,6 +233,11 @@ def get_centrality_class(self, dNchdEta: float) -> int: This function determines the index of the centrality bin for a given multiplicity value based on the predefined centrality classes. + In the case that the multiplicity input exceeds the largest or smallest + value of the multiplicity used to determine the centrality classes, the + function returns the index of the most central or most peripheral bin, + respectively. + Parameters ---------- dNchdEta : float @@ -229,17 +247,6 @@ def get_centrality_class(self, dNchdEta: float) -> int: ------- int Index of the centrality bin. - - Examples - -------- - .. highlight:: python - .. code-block:: python - :linenos: - - >>> centrality_obj = CentralityClasses(events_multiplicity=[10, 15, 20, 25], - ... centrality_bins=[0, 25, 50, 75, 100]) - >>> centrality_obj.get_centrality_class(18) - 1 """ # check if the multiplicity is in the most central bin if dNchdEta >= self.dNchdetaMin_[0]: @@ -271,16 +278,6 @@ def output_centrality_classes(self, fname: str) -> None: TypeError If :code:`fname` is not a string. - Examples - -------- - .. highlight:: python - .. code-block:: python - :linenos: - - >>> centrality_obj = CentralityClasses(events_multiplicity=[10, 15, 20, 25], - ... centrality_bins=[0, 25, 50, 75, 100]) - >>> centrality_obj.output_centrality_classes('centrality_output.txt') - Notes ----- This function writes the centrality class information, including minimum, diff --git a/src/sparkx/EventCharacteristics.py b/src/sparkx/EventCharacteristics.py index b6f0ce0f..9cbc163d 100644 --- a/src/sparkx/EventCharacteristics.py +++ b/src/sparkx/EventCharacteristics.py @@ -351,14 +351,14 @@ def generate_eBQS_densities_Milne_from_OSCAR_IC( eta_range : list, tuple A list containing the minimum and maximum values of spacetime rapidity (eta) and the number of grid points. - + output_filename : str The name of the output file where the densities will be saved. - kernel : str + kernel : str The type of kernel to use for smearing the particle data. Supported values are 'gaussian' and 'covariant'. The default is "gaussian". - + IC_info : str A string containing info about the initial condition, e.g., collision energy or centrality. @@ -405,8 +405,13 @@ def generate_eBQS_densities_Milne_from_OSCAR_IC( ) if not isinstance(self.event_data_, list): raise TypeError("The input is not a list.") - if not isinstance(kernel, str) or kernel not in {"gaussian", "covariant"}: - raise TypeError("Kernel must be a string and must be either 'covariant' or 'gaussian'.") + if not isinstance(kernel, str) or kernel not in { + "gaussian", + "covariant", + }: + raise TypeError( + "Kernel must be a string and must be either 'covariant' or 'gaussian'." + ) energy_density = Lattice3D( x_min, @@ -608,8 +613,8 @@ def generate_eBQS_densities_Minkowski_from_OSCAR_IC( output_filename : str The name of the output file where the densities will be saved. - - kernel : str + + kernel : str The type of kernel to use for smearing the particle data. Supported values are 'gaussian' and 'covariant'. The default is "gaussian". @@ -651,8 +656,13 @@ def generate_eBQS_densities_Minkowski_from_OSCAR_IC( ) if not isinstance(self.event_data_, list): raise TypeError("The input is not a list.") - if not isinstance(kernel, str) or kernel not in {"gaussian", "covariant"}: - raise TypeError("Kernel must be a string and must be either 'covariant' or 'gaussian'.") + if not isinstance(kernel, str) or kernel not in { + "gaussian", + "covariant", + }: + raise TypeError( + "Kernel must be a string and must be either 'covariant' or 'gaussian'." + ) energy_density = Lattice3D( x_min, diff --git a/src/sparkx/Filter.py b/src/sparkx/Filter.py index 0db90974..e3bd9305 100644 --- a/src/sparkx/Filter.py +++ b/src/sparkx/Filter.py @@ -319,7 +319,7 @@ def lower_event_energy_cut( List with lists containing particle objects for the events minimum_event_energy : int or float - The minimum event energy threshold. Should be a positive integer or + The minimum event energy threshold. Should be a positive integer or float. Returns @@ -680,7 +680,7 @@ def pseudorapidity_cut( cut_value : float If a single value is passed, the cut is applied symmetrically around 0. - For example, if :code:`cut_value = 1`, only particles with + For example, if :code:`cut_value = 1`, only particles with pseudo-rapidity in [-1.0, 1.0] are kept. cut_value : tuple @@ -752,7 +752,7 @@ def spacetime_rapidity_cut( cut_value : float If a single value is passed, the cut is applied symmetrically around 0. - For example, if :code:`cut_value = 1`, only particles with spacetime + For example, if :code:`cut_value = 1`, only particles with spacetime rapidity in [-1.0, 1.0] are kept. cut_value : tuple @@ -820,9 +820,9 @@ def multiplicity_cut( particle_list: List with lists containing particle objects for the events - cut_value_tuple : 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 + not in this range, the event is discarded. The range is inclusive on the lower bound and exclusive on the upper bound. Returns diff --git a/src/sparkx/Histogram.py b/src/sparkx/Histogram.py index e652b01b..d9a47a58 100755 --- a/src/sparkx/Histogram.py +++ b/src/sparkx/Histogram.py @@ -472,7 +472,7 @@ def add_value( """ Add value(s) to the latest histogram. - Different cases, if there is just one number added or a whole + Different cases, if there is just one number added or a whole list/array of numbers. Parameters diff --git a/src/sparkx/Jackknife.py b/src/sparkx/Jackknife.py index c37266df..0fb0b57e 100644 --- a/src/sparkx/Jackknife.py +++ b/src/sparkx/Jackknife.py @@ -325,9 +325,9 @@ def compute_jackknife_estimates( ) -> float: """ Compute the Jackknife uncertainty estimates for a function applied to - a data array. The default function is :code:`np.mean`, but it can be - changed to any other function that accepts a numpy array as input. - Multiple other arguments can be passed to the function as args and + a data array. The default function is :code:`np.mean`, but it can be + changed to any other function that accepts a numpy array as input. + Multiple other arguments can be passed to the function as args and kwargs. Parameters diff --git a/src/sparkx/Jetscape.py b/src/sparkx/Jetscape.py index 9a2d89bf..8982652f 100644 --- a/src/sparkx/Jetscape.py +++ b/src/sparkx/Jetscape.py @@ -85,6 +85,12 @@ class Jetscape(BaseStorer): num_events_ : int Number of events contained in the Jetscape object (updated when filters are applied) + particle_type_ : str + The type of particles in the Jetscape file (e.g. 'hadron' or 'parton') + sigmaGen_ : Tuple[float, float] + The value of sigmaGen and the uncertainty + last_line_: str + The last line of the Jetscape file Methods ------- @@ -222,7 +228,7 @@ def _particle_as_list(self, particle: Particle) -> List[Union[int, float]]: particle_list[6] = float(particle.pz) return particle_list - + def _update_after_merge(self, other: BaseStorer) -> None: """ Updates the current instance after merging with another Jetscape instance. @@ -245,16 +251,26 @@ def _update_after_merge(self, other: BaseStorer) -> None: if not isinstance(other, Jetscape): raise TypeError("Can only add Jetscape objects to Jetscape.") if self.particle_type_ != other.particle_type_: - raise TypeError("particle_types of the merged instances do not match.") - if self.particle_type_defining_string_ != other.particle_type_defining_string_: - raise TypeError("particle_type_defining_string of the merged instances do not match.") - - self.sigmaGen_ = ((self.sigmaGen_[0] + other.sigmaGen_[0])/2.0, 0.5*np.sqrt(self.sigmaGen_[1]**2 + other.sigmaGen_[1]**2)) - + raise TypeError( + "particle_types of the merged instances do not match." + ) + if ( + self.particle_type_defining_string_ + != other.particle_type_defining_string_ + ): + raise TypeError( + "particle_type_defining_string of the merged instances do not match." + ) + + self.sigmaGen_ = ( + (self.sigmaGen_[0] + other.sigmaGen_[0]) / 2.0, + 0.5 * np.sqrt(self.sigmaGen_[1] ** 2 + other.sigmaGen_[1] ** 2), + ) + # PUBLIC CLASS METHODS def participants(self) -> "Jetscape": """ - Raises an error because participants are not defined for Jetscape + Raises an error because participants are not defined for Jetscape events. Returns @@ -268,7 +284,7 @@ def participants(self) -> "Jetscape": def spectators(self) -> "Jetscape": """ - Raises an error because spectators are not defined for Jetscape + Raises an error because spectators are not defined for Jetscape events. Returns @@ -284,7 +300,7 @@ def spacetime_cut( self, dim: str, cut_value_tuple: Tuple[float, float] ) -> "Jetscape": """ - Raises an error because spacetime cuts are not possible for Jetscape + Raises an error because spacetime cuts are not possible for Jetscape events. Parameters @@ -307,7 +323,7 @@ def spacetime_rapidity_cut( self, cut_value: Union[float, Tuple[float, float]] ) -> "Jetscape": """ - Raises an error because spacetime rapidity cuts are not possible for + Raises an error because spacetime rapidity cuts are not possible for Jetscape events. Parameters @@ -325,7 +341,7 @@ def spacetime_rapidity_cut( Raises ------ NotImplementedError - Always, because spacetime rapidity cuts are not possible for + Always, because spacetime rapidity cuts are not possible for Jetscape events. """ raise NotImplementedError( @@ -362,7 +378,7 @@ def print_particle_lists_to_file(self, output_file: str) -> None: if self.num_output_per_event_ is None: raise ValueError("The number of output per event is empty.") if self.num_events_ == 1 and self.particle_list_ == [[]]: - warnings.warn("The number of events is zero.") + warnings.warn("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: f_out.write(header_file) diff --git a/src/sparkx/MultiParticlePtCorrelations.py b/src/sparkx/MultiParticlePtCorrelations.py index 016cc6d2..1d6bee4f 100644 --- a/src/sparkx/MultiParticlePtCorrelations.py +++ b/src/sparkx/MultiParticlePtCorrelations.py @@ -56,7 +56,7 @@ class MultiParticlePtCorrelations: Examples -------- - A demonstration of how to use the :code:`MultiParticlePtCorrelations` class + A demonstration of how to use the :code:`MultiParticlePtCorrelations` class to calculate transverse momentum correlations and cumulants. .. highlight:: python diff --git a/src/sparkx/Oscar.py b/src/sparkx/Oscar.py index 2bbe9398..947902c3 100644 --- a/src/sparkx/Oscar.py +++ b/src/sparkx/Oscar.py @@ -20,8 +20,8 @@ class Oscar(BaseStorer): Defines an Oscar object. The Oscar class contains a single .oscar file including all or only chosen - events in either the Oscar2013, Oscar2013Extended, Oscar2013Extended_IC, - Oscar2013Extended_Photons or ASCIICustom format. It's methods + events in either the Oscar2013, Oscar2013Extended, Oscar2013Extended_IC, + Oscar2013Extended_Photons or ASCII format. It's methods allow to directly act on all contained events as applying acceptance filters (e.g. un-/charged particles, spectators/participants) to keep/remove particles by their PDG codes or to apply cuts (e.g. multiplicity, pseudo-/rapidity, pT). @@ -77,9 +77,9 @@ class Oscar(BaseStorer): PATH_OSCAR_ : str Path to the Oscar file oscar_format_ : str - Input Oscar format "Oscar2013", "Oscar2013Extended", - "Oscar2013Extended_IC", "Oscar2013Extended_Photons", - "ASCIICustom" (set automatically) + Input Oscar format "Oscar2013", "Oscar2013Extended", + "Oscar2013Extended_IC", "Oscar2013Extended_Photons", + "ASCII" (set automatically) num_output_per_event_ : numpy.array Array containing the event number and the number of particles in this event as :code:`num_output_per_event_[event i][num_output in event i]` @@ -96,6 +96,8 @@ class Oscar(BaseStorer): ------- oscar_format: Get Oscar format of the input files + impact_parameters: + Get impact parameters of the events print_particle_lists_to_file: Print current particle data to file with same format @@ -178,6 +180,7 @@ def __init__(self, OSCAR_FILE: str, **kwargs: Any) -> None: raise TypeError("The loader must be an instance of OscarLoader.") self.oscar_format_: Union[str, None] = self.loader_.oscar_format() self.event_end_lines_: List[str] = self.loader_.event_end_lines() + self.impact_parameters_: List[float] = self.loader_.impact_parameter() del self.loader_ def create_loader(self, OSCAR_FILE: str) -> None: # type: ignore[override] @@ -200,7 +203,7 @@ def create_loader(self, OSCAR_FILE: str) -> None: # type: ignore[override] def _particle_as_list(self, particle: Any) -> List[Union[float, int]]: particle_list = [] - if self.oscar_format_ == 'ASCIICustom': + if self.oscar_format_ == "ASCII": for attr in self.custom_attr_list: particle_list.append(getattr(particle, attr)) return particle_list @@ -218,7 +221,11 @@ def _particle_as_list(self, particle: Any) -> List[Union[float, int]]: particle_list.append(int(particle.ID)) particle_list.append(int(particle.charge)) - if self.oscar_format_ == 'Oscar2013Extended' or self.oscar_format_ == 'Oscar2013Extended_IC' or self.oscar_format_ == 'Oscar2013Extended_Photons': + if ( + self.oscar_format_ == "Oscar2013Extended" + or self.oscar_format_ == "Oscar2013Extended_IC" + or self.oscar_format_ == "Oscar2013Extended_Photons" + ): particle_list.append(int(particle.ncoll)) particle_list.append(float(particle.form_time)) particle_list.append(float(particle.xsecfac)) @@ -227,7 +234,7 @@ def _particle_as_list(self, particle: Any) -> List[Union[float, int]]: particle_list.append(float(particle.t_last_coll)) particle_list.append(int(particle.pdg_mother1)) particle_list.append(int(particle.pdg_mother2)) - if self.oscar_format_ != 'Oscar2013Extended_Photons': + if self.oscar_format_ != "Oscar2013Extended_Photons": if not np.isnan(particle.baryon_number): particle_list.append(int(particle.baryon_number)) if not np.isnan(particle.strangeness): @@ -236,8 +243,15 @@ def _particle_as_list(self, particle: Any) -> List[Union[float, int]]: if not np.isnan(particle.weight): particle_list.append(int(particle.weight)) - elif self.oscar_format_ != 'Oscar2013' and self.oscar_format_ != 'Oscar2013Extended' and self.oscar_format_ != 'Oscar2013Extended_IC' and self.oscar_format_ != 'Oscar2013Extended_Photons': - raise TypeError('Input file not in OSCAR2013, OSCAR2013Extended or Oscar2013Extended_IC format') + elif ( + self.oscar_format_ != "Oscar2013" + and self.oscar_format_ != "Oscar2013Extended" + and self.oscar_format_ != "Oscar2013Extended_IC" + and self.oscar_format_ != "Oscar2013Extended_Photons" + ): + raise TypeError( + "Input file not in OSCAR2013, OSCAR2013Extended or Oscar2013Extended_IC format" + ) return particle_list @@ -245,7 +259,7 @@ def particle_status( self, status_list: Union[int, Tuple[int, ...], List[int], np.ndarray] ) -> "Oscar": """ - Raises an error because the method is not implemented for the Oscar + Raises an error because the method is not implemented for the Oscar class. Parameters @@ -268,7 +282,7 @@ def particle_status( def keep_quarks(self) -> "Oscar": """ - Raises an error because the method is not implemented for the Oscar + Raises an error because the method is not implemented for the Oscar class. Returns @@ -279,7 +293,7 @@ def keep_quarks(self) -> "Oscar": raise NotImplementedError( "keep_quarks is not implemented for the Oscar class." ) - + def _update_after_merge(self, other: BaseStorer) -> None: """ Updates the current instance after merging with another Oscar instance. @@ -300,7 +314,9 @@ def _update_after_merge(self, other: BaseStorer) -> None: if not isinstance(other, Oscar): raise TypeError("Can only add Oscar objects to Oscar.") if self.oscar_format_ != other.oscar_format_: - warnings.warn("Oscar format of the merged instances do not match. Taking the left-hand side Oscar format.") + warnings.warn( + "Oscar format of the merged instances do not match. Taking the left-hand side Oscar format." + ) self.event_end_lines_ = self.event_end_lines_ + other.event_end_lines_ def oscar_format(self) -> Union[str, None]: @@ -316,6 +332,17 @@ def oscar_format(self) -> Union[str, None]: """ return self.oscar_format_ + def impact_parameters(self) -> List[float]: + """ + Get the impact parameters for the loaded events. + + Returns + ------- + impact_parameters : List[float] + Impact parameters of the events + """ + return self.impact_parameters_ + def print_particle_lists_to_file(self, output_file: str) -> None: """ Prints the current Oscar data to an output file specified by @@ -333,32 +360,34 @@ def print_particle_lists_to_file(self, output_file: str) -> None: format_oscar2013_extended: str = ( "%g %g %g %g %g %.9g %.9g %.9g %.9g %d %d %d %d %g %g %d %d %g %d %d" ) - format_map: Dict[str,str] = { - 't': '%g', - 'x': '%g', - 'y': '%g', - 'z': '%g', - 'mass': '%g', - 'p0': '%.9g', - 'px': '%.9g', - 'py': '%.9g', - 'pz': '%.9g', - 'pdg': '%d', - 'ID': '%d', - 'charge': '%d', - 'ncoll': '%d', - 'form_time': '%g', - 'xsecfac': '%g', - 'proc_id_origin': '%d', - 'proc_type_origin': '%d', - 'time_last_coll': '%g', - 'pdg_mother1': '%g', - 'pdg_mother2': '%g', - 'baryon_number': '%g', - 'strangeness': '%g' + format_map: Dict[str, str] = { + "t": "%g", + "x": "%g", + "y": "%g", + "z": "%g", + "mass": "%g", + "p0": "%.9g", + "px": "%.9g", + "py": "%.9g", + "pz": "%.9g", + "pdg": "%d", + "ID": "%d", + "charge": "%d", + "ncoll": "%d", + "form_time": "%g", + "xsecfac": "%g", + "proc_id_origin": "%d", + "proc_type_origin": "%d", + "time_last_coll": "%g", + "pdg_mother1": "%g", + "pdg_mother2": "%g", + "baryon_number": "%g", + "strangeness": "%g", } - if self.oscar_format_ == 'ASCIICustom': - format_custom = ' '.join([format_map[attr] for attr in self.custom_attr_list]) + if self.oscar_format_ == "ASCII": + format_custom = " ".join( + [format_map[attr] for attr in self.custom_attr_list] + ) with open(self.PATH_OSCAR_, "r") as oscar_file: counter_line = 0 while True: @@ -433,7 +462,7 @@ def print_particle_lists_to_file(self, output_file: str) -> None: newline="\n", fmt=format_oscar2013_extended, ) - elif ( self.oscar_format_ == "ASCIICustom" ): + elif self.oscar_format_ == "ASCII": np.savetxt( f_out, particle_output, @@ -453,9 +482,9 @@ def print_particle_lists_to_file(self, output_file: str) -> None: f_out.write(self.event_end_lines_[event]) f_out.close() return - elif ( - len(particle_output[0]) > 20 - and (self.oscar_format_ == "Oscar2013Extended" or self.oscar_format_ == "Oscar2013Extended_IC") + elif len(particle_output[0]) > 20 and ( + self.oscar_format_ == "Oscar2013Extended" + or self.oscar_format_ == "Oscar2013Extended_IC" ): format_oscar2013_extended = ( format_oscar2013_extended @@ -481,13 +510,13 @@ def print_particle_lists_to_file(self, output_file: str) -> None: newline="\n", fmt=format_oscar2013_extended, ) - elif ( self.oscar_format_ == "ASCIICustom" ): - np.savetxt( - f_out, - particle_output, - delimiter=" ", - newline="\n", - fmt=format_custom, - ) + elif self.oscar_format_ == "ASCII": + np.savetxt( + f_out, + particle_output, + delimiter=" ", + newline="\n", + fmt=format_custom, + ) f_out.write(self.event_end_lines_[event]) f_out.close() diff --git a/src/sparkx/Particle.py b/src/sparkx/Particle.py index 627853aa..8d7fdf95 100755 --- a/src/sparkx/Particle.py +++ b/src/sparkx/Particle.py @@ -19,7 +19,7 @@ class Particle: The member variables of the Particle class are the quantities in the Oscar2013/Oscar2013Extended/Oscar2013Extended_IC/Oscar2013Extended_Photons/ - ASCIICustom or JETSCAPE hadron/parton output. If they are not set, + ASCII or JETSCAPE hadron/parton output. If they are not set, they stay :code:`np.nan` to throw an error if one tries to access a non existing quantity. If a particle with an unknown PDG is provided, a warning is thrown and and @@ -212,7 +212,7 @@ class Particle: * "Oscar2013Extended_Photons" - * "ASCIICustom" + * "ASCII" * "JETSCAPE" @@ -241,7 +241,7 @@ def __init__( self, input_format: Optional[str] = None, particle_array: Optional[np.ndarray] = None, - attribute_list: List[str] = [] + attribute_list: List[str] = [], ) -> None: self.data_: np.ndarray = np.array(25 * [np.nan], dtype=float) self.pdg_valid: bool = False @@ -250,18 +250,25 @@ def __init__( (input_format is None) and (particle_array is not None) ): raise ValueError("'input_format' or 'particle_array' not given") - - if attribute_list is None and input_format == "ASCIICustom": + + if attribute_list is None and input_format == "ASCII": raise ValueError("'attribute_list' not given") - if attribute_list!=[] and input_format != "ASCIICustom": - raise ValueError("'OscarCustom' format requires 'attribute_list' to be given.") + if attribute_list != [] and input_format != "ASCII": + raise ValueError( + "'OscarCustom' format requires 'attribute_list' to be given." + ) if (input_format is not None) and (particle_array is not None): - self.__initialize_from_array(input_format, particle_array, attribute_list) + self.__initialize_from_array( + input_format, particle_array, attribute_list + ) def __initialize_from_array( - self, input_format: str, particle_array: np.ndarray, attribute_list: List[str] + self, + input_format: str, + particle_array: np.ndarray, + attribute_list: List[str], ) -> None: """ Initialize instance attributes based on the provided input format and array, and optionally an atttribute list. @@ -274,7 +281,7 @@ def __initialize_from_array( - "Oscar2013Extended" - "Oscar2013Extended_IC" - "Oscar2013Extended_Photons" - - "ASCIICustom" + - "ASCII" - "JETSCAPE" particle_array : numpy.ndarray @@ -304,29 +311,29 @@ def __initialize_from_array( # second entry: index in line attribute_mapping = { "Allfields": { - "t": [0,0], - "x": [1,0], - "y": [2,0], - "z": [3,0], - "mass": [4,0], - "E": [5,0], - "px": [6,0], - "py": [7,0], - "pz_": [8,0], - "pdg": [9,0], - "ID": [11,0], - "charge": [12,0], - "ncoll": [13,0], - "form_time": [14,0], - "xsecfac": [15,0], - "proc_id_origin": [16,0], - "proc_type_origin": [17,0], - "t_last_coll": [18,0], - "pdg_mother1": [19,0], - "pdg_mother2": [20,0], - "status_": [21,0], - "baryon_number": [22,0], - "strangeness": [23,0], + "t": [0, 0], + "x": [1, 0], + "y": [2, 0], + "z": [3, 0], + "mass": [4, 0], + "E": [5, 0], + "px": [6, 0], + "py": [7, 0], + "pz_": [8, 0], + "pdg": [9, 0], + "ID": [11, 0], + "charge": [12, 0], + "ncoll": [13, 0], + "form_time": [14, 0], + "xsecfac": [15, 0], + "proc_id_origin": [16, 0], + "proc_type_origin": [17, 0], + "t_last_coll": [18, 0], + "pdg_mother1": [19, 0], + "pdg_mother2": [20, 0], + "status_": [21, 0], + "baryon_number": [22, 0], + "strangeness": [23, 0], }, "Oscar2013": { "t_": [0, 0], @@ -423,19 +430,31 @@ def __initialize_from_array( "pz_": [8, 6], }, } - if input_format == "ASCIICustom": + if input_format == "ASCII": mapping_dict = {} for attr in attribute_list: - mapping_dict[attr] = [attribute_mapping["Allfields"][attr][0], list(attribute_list).index(attr)] - attribute_mapping["ASCIICustom"] = mapping_dict - if input_format in attribute_mapping or input_format=="ASCIICustom": - if (input_format == "ASCIICustom" or len(particle_array) == len(attribute_mapping[input_format]) or (input_format in ["Oscar2013Extended","Oscar2013Extended_IC"]\ - and len(particle_array) <= len(attribute_mapping[input_format])\ - and len(particle_array) >= len(attribute_mapping[input_format])-2)): + mapping_dict[attr] = [ + attribute_mapping["Allfields"][attr][0], + list(attribute_list).index(attr), + ] + attribute_mapping["ASCII"] = mapping_dict + if input_format in attribute_mapping or input_format == "ASCII": + if ( + input_format == "ASCII" + or len(particle_array) == len(attribute_mapping[input_format]) + or ( + input_format + in ["Oscar2013Extended", "Oscar2013Extended_IC"] + and len(particle_array) + <= len(attribute_mapping[input_format]) + and len(particle_array) + >= len(attribute_mapping[input_format]) - 2 + ) + ): for attribute, index in attribute_mapping[input_format].items(): if len(particle_array) <= (index[1]): continue - if input_format == "ASCIICustom": + if input_format == "ASCII": attribute = attribute + "_" # Type casting for specific attributes. Although everything is saved as a float, we will only read in int data for int fields # to ensure similar behaving as if we were reading in data @@ -500,9 +519,9 @@ def __initialize_from_array( if not self.pdg_valid: if np.isnan(self.pdg): warnings.warn( - "No PDG code given! " - + "All properties extracted from the PDG are set to default values." - ) + "No PDG code given! " + + "All properties extracted from the PDG are set to default values." + ) else: warnings.warn( "The PDG code " diff --git a/src/sparkx/ParticleObjectStorer.py b/src/sparkx/ParticleObjectStorer.py index ed750772..43c95a82 100644 --- a/src/sparkx/ParticleObjectStorer.py +++ b/src/sparkx/ParticleObjectStorer.py @@ -16,7 +16,7 @@ class ParticleObjectStorer(BaseStorer): """ - Defines a :code:`ParticleObjectStorer` object, which saves particle object + Defines a :code:`ParticleObjectStorer` object, which saves particle object lists. This is a wrapper for a list of Particle objects. It's methods allow to @@ -115,9 +115,9 @@ def __init__( """ Initializes a new instance of the :code:`DummyStorer` class. - This method initializes a new instance of the :code:`DummyStorer` class - with the specified list of particle objects and optional arguments. - It calls the superclass's constructor with the particle_object_list and + This method initializes a new instance of the :code:`DummyStorer` class + with the specified list of particle objects and optional arguments. + It calls the superclass's constructor with the particle_object_list and kwargs parameters and then deletes the :code:`loader_` attribute. Parameters @@ -147,7 +147,7 @@ def create_loader( """ Creates a new :code:`ParticleObjectLoader` object. - This method creates a new :code:`ParticleObjectLoader` object with the + This method creates a new :code:`ParticleObjectLoader` object with the specified list of particle objects and assigns it to the :code:`loader_` attribute. @@ -170,10 +170,10 @@ def _particle_as_list(self, particle: "Particle") -> List[float]: """ Converts a :code:`Particle` object into a list. - This method takes a :code:`Particle` object and converts it into a list - of its attributes. The attributes are added to the list in the following - order: t, x, y, z, mass, E, px, py, pz, pdg, ID, charge, ncoll, - form_time, xsecfac, proc_id_origin, proc_type_origin, t_last_coll, + This method takes a :code:`Particle` object and converts it into a list + of its attributes. The attributes are added to the list in the following + order: t, x, y, z, mass, E, px, py, pz, pdg, ID, charge, ncoll, + form_time, xsecfac, proc_id_origin, proc_type_origin, t_last_coll, pdg_mother1, pdg_mother2, baryon_number, strangeness, weight, status. Parameters diff --git a/src/sparkx/loader/BaseLoader.py b/src/sparkx/loader/BaseLoader.py index 0c74c341..d9211d84 100644 --- a/src/sparkx/loader/BaseLoader.py +++ b/src/sparkx/loader/BaseLoader.py @@ -15,7 +15,7 @@ class BaseLoader(ABC): """ Abstract base class for all loader classes. - This class provides a common interface for all loader classes. + This class provides a common interface for all loader classes. Loader classes are responsible for loading data from a file. Attributes diff --git a/src/sparkx/loader/JetscapeLoader.py b/src/sparkx/loader/JetscapeLoader.py index 10d951e9..82b27460 100644 --- a/src/sparkx/loader/JetscapeLoader.py +++ b/src/sparkx/loader/JetscapeLoader.py @@ -157,7 +157,7 @@ def load( self.set_particle_list(kwargs), self.num_events_, self.num_output_per_event_, - [] + [], ) # PRIVATE CLASS METHODS @@ -423,7 +423,7 @@ def set_particle_list( particle_list: List[List[Particle]] = [] data: List[Particle] = [] num_read_lines = self.__get_num_read_lines() - cut_events=0 + cut_events = 0 with open(self.PATH_JETSCAPE_, "r") as jetscape_file: self._skip_lines(jetscape_file) @@ -432,23 +432,34 @@ def set_particle_list( if not line: raise IndexError("Index out of range of JETSCAPE file") elif "#" in line and "sigmaGen" in line: - old_data_len=len(data) + old_data_len = len(data) if "filters" in self.optional_arguments_.keys(): data = self.__apply_kwargs_filters( [data], kwargs["filters"] )[0] - if (len(data) != 0 or old_data_len==0): + if len(data) != 0 or old_data_len == 0: self.num_output_per_event_[len(particle_list)] = ( - len(particle_list)+1, + len(particle_list) + 1, len(data), ) else: - self.num_output_per_event_ = np.atleast_2d(np.delete(self.num_output_per_event_, len(particle_list), axis=0)) + self.num_output_per_event_ = np.atleast_2d( + np.delete( + self.num_output_per_event_, + len(particle_list), + axis=0, + ) + ) if self.num_output_per_event_.shape[0] == 0: self.num_output_per_event_ = np.array([]) - elif len(particle_list) < self.num_output_per_event_.shape[0]: - self.num_output_per_event_[len(particle_list):, 0] -= 1 - if (len(data) != 0 or old_data_len==0 ): + elif ( + len(particle_list) + < self.num_output_per_event_.shape[0] + ): + self.num_output_per_event_[ + len(particle_list) :, 0 + ] -= 1 + if len(data) != 0 or old_data_len == 0: particle_list.append(data) else: cut_events = cut_events + 1 @@ -477,23 +488,36 @@ def set_particle_list( if int(line_list[2]) == first_event_header: continue else: - old_data_len=len(data) + old_data_len = len(data) if "filters" in self.optional_arguments_.keys(): data = self.__apply_kwargs_filters( [data], kwargs["filters"] )[0] - if (len(data) != 0 or old_data_len==0): - self.num_output_per_event_[len(particle_list)] = ( - len(particle_list)+1, + if len(data) != 0 or old_data_len == 0: + self.num_output_per_event_[ + len(particle_list) + ] = ( + len(particle_list) + 1, len(data), ) else: - self.num_output_per_event_ = np.atleast_2d(np.delete(self.num_output_per_event_, len(particle_list), axis=0)) + self.num_output_per_event_ = np.atleast_2d( + np.delete( + self.num_output_per_event_, + len(particle_list), + axis=0, + ) + ) if self.num_output_per_event_.shape[0] == 0: self.num_output_per_event_ = np.array([]) - elif len(particle_list) < self.num_output_per_event_.shape[0]: - self.num_output_per_event_[len(particle_list):, 0] -= 1 - if (len(data) != 0 or old_data_len==0 ): + elif ( + len(particle_list) + < self.num_output_per_event_.shape[0] + ): + self.num_output_per_event_[ + len(particle_list) :, 0 + ] -= 1 + if len(data) != 0 or old_data_len == 0: particle_list.append(data) else: cut_events = cut_events + 1 @@ -504,8 +528,8 @@ def set_particle_list( ) particle = Particle("JETSCAPE", np.asarray(line_list)) data.append(particle) - - self.num_events_=self.num_events_-cut_events + + self.num_events_ = self.num_events_ - cut_events # Correct num_output_per_event and num_events if not kwargs or "events" not in self.optional_arguments_.keys(): if len(particle_list) != self.num_events_: diff --git a/src/sparkx/loader/OscarLoader.py b/src/sparkx/loader/OscarLoader.py index 50dd1ddb..85fa7aee 100644 --- a/src/sparkx/loader/OscarLoader.py +++ b/src/sparkx/loader/OscarLoader.py @@ -49,6 +49,8 @@ class OscarLoader(BaseLoader): Determines and sets the format of the OSCAR data file. oscar_format() Returns the OSCAR format of the data file. + impact_parameter() + Returns the impact parameter of the events in the OSCAR data file. event_end_lines() Returns the event end lines in the OSCAR data file. __get_num_read_lines() @@ -167,7 +169,7 @@ def load( self.set_particle_list(kwargs), self.num_events_, self.num_output_per_event_, - self.custom_attr_list + self.custom_attr_list, ) def _get_num_skip_lines(self) -> int: @@ -247,28 +249,28 @@ def set_num_events(self) -> None: def _set_custom_attr_list(self, header_line: List[str]) -> List[str]: self.custom_attr_list = [] attr_map = { - 't': 't', - 'x': 'x', - 'y': 'y', - 'z': 'z', - 'mass': 'mass', - 'p0': 'E', - 'px': 'px', - 'py': 'py', - 'pz': 'pz', - 'pdg': 'pdg', - 'ID': 'ID', - 'charge': 'charge', - 'ncoll': 'ncoll', - 'form_time': 'form_time', - 'xsecfac': 'xsecfac', - 'proc_id_origin': 'proc_id_origin', - 'proc_type_origin': 'proc_type_origin', - 'time_last_coll': 'time_last_coll', - 'pdg_mother1': 'pdg_mother1', - 'pdg_mother2': 'pdg_mother2', - 'baryon_number': 'baryon_number', - 'strangeness': 'strangeness' + "t": "t", + "x": "x", + "y": "y", + "z": "z", + "mass": "mass", + "p0": "E", + "px": "px", + "py": "py", + "pz": "pz", + "pdg": "pdg", + "ID": "ID", + "charge": "charge", + "ncoll": "ncoll", + "form_time": "form_time", + "xsecfac": "xsecfac", + "proc_id_origin": "proc_id_origin", + "proc_type_origin": "proc_type_origin", + "time_last_coll": "time_last_coll", + "pdg_mother1": "pdg_mother1", + "pdg_mother2": "pdg_mother2", + "baryon_number": "baryon_number", + "strangeness": "strangeness", } for i in range(0, len(header_line)): attr_name = attr_map.get(header_line[i]) @@ -316,16 +318,14 @@ def set_oscar_format(self) -> None: or first_line_list[0] == "#!OSCAR2013Extended" ): self.oscar_format_ = "Oscar2013Extended" - elif ( - first_line_list[0] == "#!ASCIICustom" - ): - self.oscar_format_ = "ASCIICustom" - value_line=first_line_list[2:] + elif first_line_list[0] == "#!ASCII": + self.oscar_format_ = "ASCII" + value_line = first_line_list[2:] self.custom_attr_list = self._set_custom_attr_list(value_line) else: raise TypeError( "Input file must follow the Oscar2013, " - + "Oscar2013Extended, Oscar2013Extended_IC, Oscar2013Extended_Photons or ASCIICustom format. " + + "Oscar2013Extended, Oscar2013Extended_IC, Oscar2013Extended_Photons or ASCII format. " ) def oscar_format(self) -> Optional[str]: @@ -356,6 +356,35 @@ def event_end_lines(self) -> List[str]: """ return self.event_end_lines_ + def impact_parameter(self) -> List[float]: + """ + Returns the impact parameter of the events. + + This method extracts the impact parameter of the collision from the + last line of each event in the OSCAR data file. + + Returns + ------- + List[float] + The impact parameter of the collisions. + """ + impact_parameters: List[float] = [] + for line in self.event_end_lines_: + line_split = line.split(" ") + line_split = list(filter(None, line_split)) + impact_parameter = float(line_split[-3]) + impact_parameters.append(impact_parameter) + + # update the list with the num_output_per_event_ for the case + # that only certain events are loaded from the file + if self.num_output_per_event_.shape[0] == 0: + impact_parameters = [] + else: + idx_list = self.num_output_per_event_[:, 0] + impact_parameters = [impact_parameters[i] for i in idx_list] + + return impact_parameters + def __get_num_read_lines(self) -> int: """ Calculates the number of lines to read from the OSCAR data file. @@ -533,8 +562,6 @@ def __apply_kwargs_filters( return event - # PUBLIC CLASS METHODS - def set_particle_list(self, kwargs: Dict[str, Any]) -> List[List[Particle]]: """ Sets the list of particles from the OSCAR data file. @@ -571,7 +598,7 @@ def set_particle_list(self, kwargs: Dict[str, Any]) -> List[List[Particle]]: particle_list: List[List[Particle]] = [] data: List[Particle] = [] num_read_lines = self.__get_num_read_lines() - cut_events=0 + cut_events = 0 with open(self.PATH_OSCAR_, "r") as oscar_file: self._skip_lines(oscar_file) for i in range(0, num_read_lines): @@ -587,26 +614,39 @@ def set_particle_list(self, kwargs: Dict[str, Any]) -> List[List[Particle]]: "First line of the event is not a comment " + 'line or does not contain "out"' ) - elif "event" in line and ("out" in line or "in " in line or " start" in line): + elif "event" in line and ( + "out" in line or "in " in line or " start" in line + ): continue elif "#" in line and "end" in line: - old_data_len=len(data) + old_data_len = len(data) if "filters" in self.optional_arguments_.keys(): data = self.__apply_kwargs_filters( [data], kwargs["filters"] )[0] - if (len(data) != 0 or old_data_len==0): + if len(data) != 0 or old_data_len == 0: self.num_output_per_event_[len(particle_list)] = ( len(particle_list), len(data), ) else: - self.num_output_per_event_ = np.atleast_2d(np.delete(self.num_output_per_event_, len(particle_list), axis=0)) + self.num_output_per_event_ = np.atleast_2d( + np.delete( + self.num_output_per_event_, + len(particle_list), + axis=0, + ) + ) if self.num_output_per_event_.shape[0] == 0: self.num_output_per_event_ = np.array([]) - elif len(particle_list) < self.num_output_per_event_.shape[0]: - self.num_output_per_event_[len(particle_list):, 0] -= 1 - if (len(data) != 0 or old_data_len==0 ): + elif ( + len(particle_list) + < self.num_output_per_event_.shape[0] + ): + self.num_output_per_event_[ + len(particle_list) :, 0 + ] -= 1 + if len(data) != 0 or old_data_len == 0: particle_list.append(data) else: cut_events = cut_events + 1 @@ -615,12 +655,14 @@ def set_particle_list(self, kwargs: Dict[str, Any]) -> List[List[Particle]]: raise ValueError("Comment line unexpectedly found: " + line) else: line_list = np.asarray(line.replace("\n", "").split(" ")) - if(self.oscar_format_ == "ASCIICustom"): - particle = Particle(self.oscar_format_, line_list, self.custom_attr_list) + if self.oscar_format_ == "ASCII": + particle = Particle( + self.oscar_format_, line_list, self.custom_attr_list + ) else: particle = Particle(self.oscar_format_, line_list) data.append(particle) - self.num_events_=self.num_events_-cut_events + self.num_events_ = self.num_events_ - cut_events # Correct num_output_per_event and num_events if not kwargs or "events" not in self.optional_arguments_.keys(): if len(particle_list) != self.num_events_: @@ -703,7 +745,7 @@ def set_num_output_per_event_and_event_footers(self) -> None: elif "#" in line and " end" in line: self.event_end_lines_.append(line) event_output.append([event, line_counter - 1]) - elif "#" in line and ( " in " in line or " start" in line): + elif "#" in line and (" in " in line or " start" in line): line_str = line.replace("\n", "").split(" ") event = int(line_str[2]) line_counter = 0 diff --git a/src/sparkx/loader/ParticleObjectLoader.py b/src/sparkx/loader/ParticleObjectLoader.py index 969c5e93..236de5ce 100644 --- a/src/sparkx/loader/ParticleObjectLoader.py +++ b/src/sparkx/loader/ParticleObjectLoader.py @@ -144,7 +144,7 @@ def load( self.set_particle_list(kwargs), self.num_events_, self.num_output_per_event_, - [] + [], ) def set_num_output_per_event(self) -> List[int]: diff --git a/tests/test_BaseStorer.py b/tests/test_BaseStorer.py index 7fb281f3..ba83a814 100644 --- a/tests/test_BaseStorer.py +++ b/tests/test_BaseStorer.py @@ -42,10 +42,11 @@ def print_particle_lists_to_file(self, output_file) -> None: with open(output_file, "w") as file: for event in self.particle_list(): file.write(str(event) + "\n") - + def _update_after_merge(self, other: "ConcreteStorer") -> None: pass + class ConcreteStorer2(ConcreteStorer): pass @@ -111,6 +112,7 @@ def test_print_particle_lists_to_file(concrete_storer, tmp_path): content = f.read() assert content == "[1, 2, 3]\n[4, 5, 6]\n" + def test_add_storer(): storer1 = ConcreteStorer(path="dummy_path") storer2 = ConcreteStorer(path="dummy_path") @@ -120,16 +122,13 @@ def test_add_storer(): storer1.particle_list_ = [ [[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]], - [[7, 8, 9], [10, 11, 12], [10, 11, 12]] + [[7, 8, 9], [10, 11, 12], [10, 11, 12]], ] storer1.num_output_per_event_ = np.array([[1, 2], [2, 2], [3, 3]]) storer1.num_events_ = 3 # Mock data for storer2 - storer2.particle_list_ = [ - [[13, 14, 15]], - [[19, 20, 21], [22, 23, 24]] - ] + storer2.particle_list_ = [[[13, 14, 15]], [[19, 20, 21], [22, 23, 24]]] storer2.num_output_per_event_ = np.array([[1, 1], [2, 2]]) storer2.num_events_ = 2 @@ -142,12 +141,14 @@ def test_add_storer(): [[7, 8, 9], [10, 11, 12]], [[7, 8, 9], [10, 11, 12], [10, 11, 12]], [[13, 14, 15]], - [[19, 20, 21], [22, 23, 24]] + [[19, 20, 21], [22, 23, 24]], ] - expected_num_output_per_event = np.array([[1, 2], [2, 2],[3, 3], [4, 1], [5, 2]]) + expected_num_output_per_event = np.array( + [[1, 2], [2, 2], [3, 3], [4, 1], [5, 2]] + ) expected_num_events = 5 - # Add the two storers + # Add the two storers combined_storer2 = storer2 + storer1 # Expected combined data @@ -156,21 +157,27 @@ def test_add_storer(): [[19, 20, 21], [22, 23, 24]], [[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]], - [[7, 8, 9], [10, 11, 12], [10, 11, 12]] + [[7, 8, 9], [10, 11, 12], [10, 11, 12]], ] - expected_num_output_per_event2 = np.array([[1, 1], [2, 2],[3, 2], [4, 2], [5, 3]]) + expected_num_output_per_event2 = np.array( + [[1, 1], [2, 2], [3, 2], [4, 2], [5, 3]] + ) expected_num_events2 = 5 # Assertions assert combined_storer.particle_list_ == expected_particle_list - assert np.array_equal(combined_storer.num_output_per_event_, expected_num_output_per_event) + assert np.array_equal( + combined_storer.num_output_per_event_, expected_num_output_per_event + ) assert combined_storer.num_events_ == expected_num_events assert combined_storer2.particle_list_ == expected_particle_list2 - assert np.array_equal(combined_storer2.num_output_per_event_, expected_num_output_per_event2) + assert np.array_equal( + combined_storer2.num_output_per_event_, expected_num_output_per_event2 + ) assert combined_storer2.num_events_ == expected_num_events2 with pytest.raises(TypeError): storer1 + 1 - + with pytest.raises(TypeError): - storer1 + storer3 \ No newline at end of file + storer1 + storer3 diff --git a/tests/test_EventCharacteristics.py b/tests/test_EventCharacteristics.py index 099f2df3..c4ff8160 100644 --- a/tests/test_EventCharacteristics.py +++ b/tests/test_EventCharacteristics.py @@ -177,7 +177,7 @@ def test_eBQS_densities_Milne_from_OSCAR_IC( n_sigma_x = n_sigma_y = n_sigma_z = 2 sigma_smear = 0.3 kernel = "covariant" - eta_range = [-1, 1, 10] + eta_range = [-1, 1, 10] output_filename = os.path.join( tmp_path, "test_eBQS_densities_Milne_from_OSCAR_IC.dat" ) @@ -202,7 +202,7 @@ def test_eBQS_densities_Milne_from_OSCAR_IC( eta_range, output_filename, kernel, - IC_info=1 + IC_info=1, ) # test error if class is initialized with lattice @@ -224,7 +224,7 @@ def test_eBQS_densities_Milne_from_OSCAR_IC( sigma_smear, eta_range, output_filename, - kernel + kernel, ) # Check error if x_min, x_max, y_min, y_max, z_min, z_max are not floats with pytest.raises(TypeError): @@ -245,7 +245,7 @@ def test_eBQS_densities_Milne_from_OSCAR_IC( sigma_smear, eta_range, output_filename, - kernel + kernel, ) # Check error if Nx, Ny, Nz are not positive integers with pytest.raises(TypeError): @@ -266,7 +266,7 @@ def test_eBQS_densities_Milne_from_OSCAR_IC( sigma_smear, eta_range, output_filename, - kernel + kernel, ) with pytest.raises(TypeError): eve_char = EventCharacteristics(oscar_particle_objects_list) @@ -286,7 +286,7 @@ def test_eBQS_densities_Milne_from_OSCAR_IC( sigma_smear, eta_range, output_filename, - kernel + kernel, ) # Check error if n_sigma_x, n_sigma_y, n_sigma_z are not floats with pytest.raises(TypeError): @@ -307,7 +307,7 @@ def test_eBQS_densities_Milne_from_OSCAR_IC( sigma_smear, eta_range, output_filename, - kernel + kernel, ) with pytest.raises(TypeError): eve_char = EventCharacteristics(oscar_particle_objects_list) @@ -327,7 +327,7 @@ def test_eBQS_densities_Milne_from_OSCAR_IC( sigma_smear, eta_range, output_filename, - kernel + kernel, ) # Check error if kernel is not a string 'covariant' or 'gaussian' with pytest.raises(TypeError): @@ -348,8 +348,8 @@ def test_eBQS_densities_Milne_from_OSCAR_IC( sigma_smear, eta_range, output_filename, - 4 - ) + 4, + ) with pytest.raises(TypeError): eve_char = EventCharacteristics(oscar_particle_objects_list) eve_char.generate_eBQS_densities_Milne_from_OSCAR_IC( @@ -368,8 +368,8 @@ def test_eBQS_densities_Milne_from_OSCAR_IC( sigma_smear, eta_range, output_filename, - "invalid" - ) + "invalid", + ) # Check error if eta_range is not a list with 3 float entries with pytest.raises(ValueError): eve_char = EventCharacteristics(oscar_particle_objects_list) @@ -389,7 +389,7 @@ def test_eBQS_densities_Milne_from_OSCAR_IC( sigma_smear, [1, 2], output_filename, - kernel + kernel, ) with pytest.raises(TypeError): eve_char = EventCharacteristics(oscar_particle_objects_list) @@ -409,7 +409,7 @@ def test_eBQS_densities_Milne_from_OSCAR_IC( sigma_smear, [1, "invalid", 2], output_filename, - kernel + kernel, ) with pytest.raises(TypeError): eve_char = EventCharacteristics(oscar_particle_objects_list) @@ -429,7 +429,7 @@ def test_eBQS_densities_Milne_from_OSCAR_IC( sigma_smear, 3, output_filename, - kernel + kernel, ) # Check error if output_filename is not a string with pytest.raises(TypeError): @@ -450,7 +450,7 @@ def test_eBQS_densities_Milne_from_OSCAR_IC( sigma_smear, eta_range, 1, - kernel + kernel, ) # Check error if not enough data for computing tau @@ -472,7 +472,7 @@ def test_eBQS_densities_Milne_from_OSCAR_IC( sigma_smear, eta_range, output_filename, - kernel + kernel, ) # Generate densities @@ -576,7 +576,7 @@ def test_eBQS_densities_Minkowski_from_OSCAR_IC( sigma_smear, output_filename, kernel, - IC_info=1 + IC_info=1, ) # test error if class is initialized with lattice with pytest.raises(TypeError): @@ -596,7 +596,7 @@ def test_eBQS_densities_Minkowski_from_OSCAR_IC( n_sigma_z, sigma_smear, output_filename, - kernel + kernel, ) # Check error if x_min, x_max, y_min, y_max, z_min, z_max are not floats with pytest.raises(TypeError): @@ -616,7 +616,7 @@ def test_eBQS_densities_Minkowski_from_OSCAR_IC( n_sigma_z, sigma_smear, output_filename, - kernel + kernel, ) # Check error if Nx, Ny, Nz are not positive integers with pytest.raises(TypeError): @@ -636,7 +636,7 @@ def test_eBQS_densities_Minkowski_from_OSCAR_IC( n_sigma_z, sigma_smear, output_filename, - kernel + kernel, ) with pytest.raises(TypeError): eve_char = EventCharacteristics(oscar_particle_objects_list) @@ -655,7 +655,7 @@ def test_eBQS_densities_Minkowski_from_OSCAR_IC( n_sigma_z, sigma_smear, output_filename, - kernel + kernel, ) # Check error if n_sigma_x, n_sigma_y, n_sigma_z are not positive floats with pytest.raises(TypeError): @@ -675,7 +675,7 @@ def test_eBQS_densities_Minkowski_from_OSCAR_IC( n_sigma_z, sigma_smear, output_filename, - kernel + kernel, ) with pytest.raises(TypeError): eve_char = EventCharacteristics(oscar_particle_objects_list) @@ -694,7 +694,7 @@ def test_eBQS_densities_Minkowski_from_OSCAR_IC( n_sigma_z, sigma_smear, output_filename, - kernel + kernel, ) # Check error if kernel is not a string 'covariant' or 'gaussian' with pytest.raises(TypeError): @@ -714,8 +714,8 @@ def test_eBQS_densities_Minkowski_from_OSCAR_IC( n_sigma_z, sigma_smear, output_filename, - 4 - ) + 4, + ) with pytest.raises(TypeError): eve_char = EventCharacteristics(oscar_particle_objects_list) eve_char.generate_eBQS_densities_Minkowski_from_OSCAR_IC( @@ -733,8 +733,8 @@ def test_eBQS_densities_Minkowski_from_OSCAR_IC( n_sigma_z, sigma_smear, output_filename, - "invalid" - ) + "invalid", + ) # Check error if output_filename is not a string with pytest.raises(TypeError): eve_char = EventCharacteristics(oscar_particle_objects_list) @@ -753,5 +753,5 @@ def test_eBQS_densities_Minkowski_from_OSCAR_IC( n_sigma_z, sigma_smear, 1, - kernel + kernel, ) diff --git a/tests/test_Filter.py b/tests/test_Filter.py index b85e0882..a513db31 100644 --- a/tests/test_Filter.py +++ b/tests/test_Filter.py @@ -171,7 +171,9 @@ def particle_list_hadrons_with_up_down_quarks(): return [particle_list] -def test_keep_up_quark_hadrons(particle_nan_quantities, particle_list_hadrons_with_up_down_quarks): +def test_keep_up_quark_hadrons( + particle_nan_quantities, particle_list_hadrons_with_up_down_quarks +): return_list = keep_up(particle_nan_quantities) assert len(return_list[0]) == 0 @@ -179,7 +181,9 @@ def test_keep_up_quark_hadrons(particle_nan_quantities, particle_list_hadrons_wi assert len(return_list[0]) == 5 -def test_keep_down_quark_hadrons(particle_nan_quantities, particle_list_hadrons_with_up_down_quarks): +def test_keep_down_quark_hadrons( + particle_nan_quantities, particle_list_hadrons_with_up_down_quarks +): return_list = keep_down(particle_nan_quantities) assert len(return_list[0]) == 0 @@ -201,13 +205,16 @@ def particle_list_hadrons_with_charm_quarks(): return [particle_list] -def test_keep_charm_quark_hadrons(particle_nan_quantities, particle_list_hadrons_with_charm_quarks): +def test_keep_charm_quark_hadrons( + particle_nan_quantities, particle_list_hadrons_with_charm_quarks +): return_list = keep_charm(particle_nan_quantities) assert len(return_list[0]) == 0 return_list = keep_charm(particle_list_hadrons_with_charm_quarks) assert len(return_list[0]) == 4 + @pytest.fixture def particle_list_hadrons_with_bottom_quarks(): particle_list = [] @@ -222,7 +229,9 @@ def particle_list_hadrons_with_bottom_quarks(): return [particle_list] -def test_keep_bottom_quark_hadrons(particle_nan_quantities, particle_list_hadrons_with_bottom_quarks): +def test_keep_bottom_quark_hadrons( + particle_nan_quantities, particle_list_hadrons_with_bottom_quarks +): return_list = keep_bottom(particle_nan_quantities) assert len(return_list[0]) == 0 @@ -230,7 +239,9 @@ def test_keep_bottom_quark_hadrons(particle_nan_quantities, particle_list_hadron assert len(return_list[0]) == 4 -def test_keep_top_quark(particle_nan_quantities, particle_list_hadrons_with_bottom_quarks): +def test_keep_top_quark( + particle_nan_quantities, particle_list_hadrons_with_bottom_quarks +): return_list = keep_top(particle_nan_quantities) assert len(return_list[0]) == 0 @@ -252,7 +263,9 @@ def particle_list_hadrons_with_photons(): return [particle_list] -def test_remove_photons(particle_nan_quantities, particle_list_hadrons_with_photons): +def test_remove_photons( + particle_nan_quantities, particle_list_hadrons_with_photons +): return_list = remove_photons(particle_nan_quantities) assert len(return_list[0]) == 0 @@ -770,23 +783,24 @@ def particle_list_multiplicity(): def test_multiplicity_cut(particle_list_multiplicity): # Test cases for valid input - assert multiplicity_cut(particle_list_multiplicity, (7,None)) == [ + assert multiplicity_cut(particle_list_multiplicity, (7, None)) == [ particle_list_multiplicity[1] ] - assert multiplicity_cut(particle_list_multiplicity, (5,10)) == [ + assert multiplicity_cut(particle_list_multiplicity, (5, 10)) == [ particle_list_multiplicity[0] ] assert multiplicity_cut(particle_list_multiplicity, (11, None)) == [[]] - assert multiplicity_cut(particle_list_multiplicity, (10, 5))== [ + assert multiplicity_cut(particle_list_multiplicity, (10, 5)) == [ particle_list_multiplicity[0] ] # Test cases for invalid input with pytest.raises(TypeError): - multiplicity_cut(particle_list_multiplicity, cut_value=(-3.5,4)) + multiplicity_cut(particle_list_multiplicity, cut_value=(-3.5, 4)) with pytest.raises(TypeError): - multiplicity_cut(particle_list_multiplicity, cut_value=(0,'a')) + multiplicity_cut(particle_list_multiplicity, cut_value=(0, "a")) + @pytest.fixture def particle_list_status(): diff --git a/tests/test_Jetscape.py b/tests/test_Jetscape.py index e26be67d..7badfb29 100644 --- a/tests/test_Jetscape.py +++ b/tests/test_Jetscape.py @@ -370,9 +370,9 @@ def test_filter_in_Jetscape_constructor(jetscape_file_path): np.array([]), ) jetscape_empty_multiplicity = Jetscape( - jetscape_file_path, filters={"multiplicity_cut": (99999999,None)} + jetscape_file_path, filters={"multiplicity_cut": (99999999, None)} ) - print( jetscape_empty_multiplicity.num_output_per_event()) + print(jetscape_empty_multiplicity.num_output_per_event()) assert np.array_equal( jetscape_empty_multiplicity.num_output_per_event(), np.array([]) ) @@ -478,14 +478,14 @@ def test_Jetscape_print_with_empty_events( assert filecmp.cmp(jetscape_file_no_hadrons, output_path) os.remove(output_path) -def test_Jetscape_print_with_no_events( - jetscape_file_path, output_path -): - jetscape = Jetscape(jetscape_file_path).multiplicity_cut((100000000,None)) + +def test_Jetscape_print_with_no_events(jetscape_file_path, output_path): + jetscape = Jetscape(jetscape_file_path).multiplicity_cut((100000000, None)) with pytest.warns(UserWarning): jetscape.print_particle_lists_to_file(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) @@ -514,17 +514,18 @@ def test_Jetscape_read_parton_file(jetscape_file_path_partons): with pytest.raises(TypeError): Jetscape(jetscape_file_path_partons, particletype=1) + def test_update_after_merge_warning(jetscape_file_path): # Create two Jetscape instances with different particle types and defining strings jetscape1 = Jetscape(jetscape_file_path) jetscape1.particle_type_ = "parton" jetscape1.particle_type_defining_string_ = "parton" - jetscape1.sigmaGen_ = (1,2) + jetscape1.sigmaGen_ = (1, 2) jetscape2 = Jetscape(jetscape_file_path) jetscape2.particle_type_ = "hadron" jetscape2.particle_type_defining_string_ = "hadron" - jetscape2.sigmaGen_ = (2,2) + jetscape2.sigmaGen_ = (2, 2) # Check if a UserWarning is issued when merging with pytest.raises(TypeError): @@ -538,4 +539,4 @@ def test_update_after_merge_warning(jetscape_file_path): jetscape1._update_after_merge(jetscape2) - assert jetscape1.sigmaGen_ == (1.5,np.sqrt(2)) \ No newline at end of file + assert jetscape1.sigmaGen_ == (1.5, np.sqrt(2)) diff --git a/tests/test_Oscar.py b/tests/test_Oscar.py index dc5bb292..70ea2210 100644 --- a/tests/test_Oscar.py +++ b/tests/test_Oscar.py @@ -238,6 +238,27 @@ def test_oscar_format(tmp_path): assert oscar.oscar_format() == "Oscar2013Extended" +def test_oscar_impact_parameter(oscar_extended_file_path): + oscar = Oscar(oscar_extended_file_path) + impact_parameters = oscar.impact_parameters() + assert impact_parameters == [0.0, 1.0, 2.0, 3.0, 4.0] + + oscar1 = Oscar(oscar_extended_file_path, events=(1)) + impact_parameters = oscar1.impact_parameters() + assert impact_parameters == [1.0] + + oscar2 = Oscar(oscar_extended_file_path, events=(1, 2)) + impact_parameters = oscar2.impact_parameters() + assert impact_parameters == [1.0, 2.0] + + oscar1 = Oscar( + oscar_extended_file_path, filters={"multiplicity_cut": (50, None)} + ) + impact_parameters = oscar1.impact_parameters() + # When all events are empty we obtain an empty list + assert impact_parameters == [] + + def test_num_output_per_event(tmp_path, oscar_old_extended_file_path): num_events = 7 num_output_per_event = [1, 35, 27, 0, 9, 17, 3] @@ -599,7 +620,7 @@ def test_filter_in_oscar_constructor(tmp_path): oscar_file, filters={"uncharged_particles": True, "spectators": True} ) oscar_empty_multiplicity = Oscar( - oscar_file, filters={"multiplicity_cut": (99999999,None)} + oscar_file, filters={"multiplicity_cut": (99999999, None)} ) assert np.array_equal( @@ -608,9 +629,7 @@ def test_filter_in_oscar_constructor(tmp_path): assert np.array_equal( oscar_uncharged.num_output_per_event(), np.array([[0, 18], [1, 20]]) ) - assert np.array_equal( - oscar_empty.num_output_per_event(), np.array([]) - ) + assert np.array_equal(oscar_empty.num_output_per_event(), np.array([])) assert np.array_equal( oscar_charged_participants.num_output_per_event(), np.array([[0, 12], [1, 10]]), @@ -622,7 +641,14 @@ def test_filter_in_oscar_constructor(tmp_path): assert np.array_equal( oscar_empty_multiplicity.num_output_per_event(), np.array([]) ) - del oscar_file, oscar_charged, oscar_uncharged, oscar_empty, oscar_uncharged_spectators, oscar_empty_multiplicity + del ( + oscar_file, + oscar_charged, + oscar_uncharged, + oscar_empty, + oscar_uncharged_spectators, + oscar_empty_multiplicity, + ) # Create a particle list to be loaded into the Oscar object oscar_file = str(tmp_path / "particle_lists.oscar") @@ -677,9 +703,7 @@ def test_filter_in_oscar_constructor(tmp_path): assert np.array_equal( oscar_spectators.num_output_per_event(), np.array([[0, 20], [1, 22]]) ) - assert np.array_equal( - oscar_empty.num_output_per_event(), np.array([]) - ) + assert np.array_equal(oscar_empty.num_output_per_event(), np.array([])) assert np.array_equal( oscar_spectators_strange.num_output_per_event(), np.array([[0, 15], [1, 22]]), @@ -704,7 +728,7 @@ 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)) + 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) @@ -728,12 +752,16 @@ def test_old_extended_oscar_print(oscar_old_extended_file_path, output_path): def test_custom_oscar_print(oscar_custom_file_path, output_path): - with pytest.warns(UserWarning, match="No PDG code given! All properties extracted from the PDG are set to default values."): + with pytest.warns( + UserWarning, + match="No PDG code given! All properties extracted from the PDG are set to default values.", + ): oscar = Oscar(oscar_custom_file_path) oscar.print_particle_lists_to_file(output_path) assert filecmp.cmp(oscar_custom_file_path, output_path) os.remove(output_path) + def test_update_after_merge_warning(oscar_file_path): # Create two Oscar instances with different formats oscar1 = Oscar(oscar_file_path) @@ -745,9 +773,10 @@ def test_update_after_merge_warning(oscar_file_path): oscar2.event_end_lines_ = [40, 50, 60] # Check if a UserWarning is issued when merging - with pytest.warns(UserWarning, match="Oscar format of the merged instances do not match"): + with pytest.warns( + UserWarning, match="Oscar format of the merged instances do not match" + ): oscar1._update_after_merge(oscar2) # Check if the event_end_lines_ are updated correctly assert oscar1.event_end_lines_ == [10, 20, 30, 40, 50, 60] - diff --git a/tests/test_Particle.py b/tests/test_Particle.py index 21c5188f..026ece78 100644 --- a/tests/test_Particle.py +++ b/tests/test_Particle.py @@ -827,6 +827,7 @@ def test_is_quark_from_pdg_valid_values(): assert result == expected_result assert result == True + def test_is_quark_from_pdg_invalid_values(): p = Particle() p.pdg_valid = False diff --git a/tests/test_files/particle_lists_extended.oscar b/tests/test_files/particle_lists_extended.oscar index c5346399..82b69b2a 100644 --- a/tests/test_files/particle_lists_extended.oscar +++ b/tests/test_files/particle_lists_extended.oscar @@ -68,7 +68,7 @@ 200 -0.418609 -0.176277 -67.3449 0.938 0.969357088 -0.0320687583 0.0551601951 -0.236089203 2212 29 1 0 -5.76975 1 0 0 0 0 0 1 0 200 2.42947 3.0112 -65.1327 0.938 1.04865056 0.0394574081 -0.0777718654 -0.46067196 2212 30 1 0 -5.76975 1 0 0 0 0 0 1 0 200 0.378057 0.128352 -64.5926 0.938 1.00775805 -0.150150434 -0.111678218 -0.317356455 2212 31 1 0 -5.76975 1 0 0 0 0 0 1 0 -# event 1 end 0 impact 0.000 scattering_projectile_target yes +# event 1 end 0 impact 1.000 scattering_projectile_target yes # event 2 out 32 200 1.43681 -0.57951 68.1254 0.938 0.984674176 0.101050592 0.0217277201 0.281168839 2112 0 0 0 -5.76975 1 0 0 0 0 0 1 0 200 -66.8826 25.0482 48.7983 0.938 1.05862791 -0.372121994 0.142474767 0.286487728 2112 1 0 1 -5.76975 1 5 1 16.2822 0 0 1 0 @@ -102,7 +102,7 @@ 200 52.4179 -5.45205 -52.4466 0.938 1.0167732 0.269232926 -0.037541421 -0.282998271 2212 29 1 1 -5.76975 1 1 1 8.285 0 0 1 0 200 2.07186 1.42018 -65.5257 0.938 0.997005975 0.0217456011 -0.172661831 -0.289641045 2212 30 1 0 -5.76975 1 0 0 0 0 0 1 0 200 -0.328325 0.0840569 -67.7737 0.938 0.977507113 -0.0151826403 -0.112229299 -0.250699477 2212 31 1 0 -5.76975 1 0 0 0 0 0 1 0 -# event 2 end 0 impact 0.000 scattering_projectile_target yes +# event 2 end 0 impact 2.000 scattering_projectile_target yes # event 3 out 32 200 32.5875 -43.6152 90.6772 0.938 1.13424733 0.200234283 -0.258609143 0.547449133 2112 0 0 3 -5.76975 1 13 1 13.2923 0 0 1 0 200 -1.38162 -0.233909 66.1332 0.938 1.00757408 0.108197762 -0.115310548 0.332202126 2112 1 0 0 -5.76975 1 0 0 0 0 0 1 0 @@ -136,7 +136,7 @@ 200 2.22264 -3.63981 -65.796 0.938 1.03388647 0.0402788199 0.00187661296 -0.432956507 2212 29 1 0 -5.76975 1 0 0 0 0 0 1 0 200 -0.566019 2.40523 -66.3951 0.938 0.97863307 0.0669109903 -0.012243816 -0.270650502 2212 30 1 0 -5.76975 1 0 0 0 0 0 1 0 200 49.1038 30.8507 11.5134 0.938 0.987926588 0.259317497 0.161110396 0.0543398555 2212 31 1 1 -5.76975 1 12 1 12.9886 0 0 1 0 -# event 3 end 0 impact 0.000 scattering_projectile_target yes +# event 3 end 0 impact 3.000 scattering_projectile_target yes # event 4 out 32 200 90.1624 44.9239 33.299 0.938 1.1173779 0.51704332 0.250424132 0.196579062 2112 0 0 2 -5.76975 1 12 1 10.8119 0 0 1 0 200 -0.498419 3.71458 65.8835 0.938 0.990781194 -0.0249362272 -0.070860377 0.310097348 2112 1 0 0 -5.76975 1 0 0 0 0 0 1 0 @@ -170,4 +170,4 @@ 200 59.8822 15.2357 10.3679 0.938 0.992192205 0.309457517 0.0708851794 0.0617471318 2212 29 1 2 -5.76975 1 15 1 14.9758 0 0 1 0 200 88.6663 -9.57923 6.69219 0.938 1.05750388 0.483828783 -0.058722314 0.0305262982 2212 30 1 1 -5.76975 1 6 1 6.98856 0 0 1 0 200 47.3753 5.87873 -38.0641 0.938 0.992125203 0.245840815 0.0396484427 -0.20605512 2212 31 1 1 -5.76975 1 9 1 8.99191 0 0 1 0 -# event 4 end 0 impact 0.000 scattering_projectile_target yes +# event 4 end 0 impact 4.000 scattering_projectile_target yes diff --git a/tests/test_files/test_custom_oscar.dat b/tests/test_files/test_custom_oscar.dat index f3672741..f9ef35b9 100644 --- a/tests/test_files/test_custom_oscar.dat +++ b/tests/test_files/test_custom_oscar.dat @@ -1,4 +1,4 @@ -#!ASCIICustom particle_lists px py z +#!ASCII particle_lists px py z # Units: GeV GeV fm # SMASH-3.1-220-ge0fbc0856 # event 0 out 485