import numpy as np
import random as rd
[docs]class GenerateFlow:
"""
Generate particle data with anisotropic flow for testing.
This class generates particle lists in JETSCAPE or OSCAR output format
to test the correct implementation of flow analysis routines.
Attributes
----------
n_:
Type of flow harmonics.
vn_:
Value of the flow harmonics.
phi_:
List containing the azimuths of the particles.
px_:
Particle momenta in x-direction.
py_:
Particle momenta in y-direction.
pz_:
Particle momenta in z-direction.
Methods
-------
generate_dummy_JETSCAPE_file:
Generate dummy flow data in JETSCAPE format.
generate_dummy_OSCAR_file:
Generate dummy flow data in OSCAR format.
Examples
--------
To use the class the GenerateFlow object has to be created with the desired
anisotropic flow harmonics and then a dummy data file can be created:
.. highlight:: python
.. code-block:: python
:linenos:
>>> from GenerateFlow import GenerateFlow
>>>
>>> flow_object = GenerateFlow(v2=0.06, v3=0.02, v4=0.03)
>>> number_events = 100
>>> event_multiplicity = 10000
>>> random_seed = 42
>>> flow_object.generate_dummy_JETSCAPE_file(path_to_output,number_events,event_multiplicity,random_seed)
"""
def __init__(self, *vn, **vn_kwargs):
if not vn and not vn_kwargs:
self.n_ = self.vn_ = None
else:
try:
vn_dictionary = {int(kw.lstrip('v')): val for kw, val in vn_kwargs.items()}
except ValueError:
raise TypeError("Input must have the form of a dictionary with 'vN' "
"where N is an integer.")
vn_dictionary.update((k, v) for k, v in enumerate(vn, start=2) if v is not None and v != 0.)
kwargs = dict(dtype=float, count=len(vn_dictionary))
self.n_ = np.fromiter(vn_dictionary.keys(), **kwargs)
self.vn_ = np.fromiter(vn_dictionary.values(), **kwargs)
self.phi_ = []
self.px_ = []
self.py_ = []
self.pz_ = []
def __distribution_function(self,phi):
"""
Calculate the distribution function value for a given angle.
Parameters
----------
phi: float
The azimuthal angle.
Returns
-------
float
The value of the distribution function at the given angle.
Notes
-----
- This method calculates the value of the distribution function
based on the harmonic components (vn) and angles (phi).
"""
f = 1. / (2. * np.pi)
f_harmonic = 1.0
for term in range(len(self.n_)):
f_harmonic += 2. * self.vn_[term] * np.cos(self.n_[term] * phi)
return f * f_harmonic
def sample_angles(self, multiplicity):
"""
Sample angles for a given multiplicity according to a given distribution.
Parameters
----------
multiplicity: int
The number of angles to sample.
Returns
-------
None
"""
f_max = (1. + 2. * self.vn_.sum()) / (2. * np.pi)
phi = []
while len(phi) < multiplicity:
random_phi = rd.uniform(0., 2. * np.pi)
random_dist_val = rd.uniform(0., f_max)
if random_dist_val <= self.__distribution_function(random_phi):
phi.append(random_phi)
self.phi_ = phi
def __thermal_distribution(self,temperature,mass):
"""
Calculate the momentum magnitude from a thermal distribution.
Parameters
----------
temperature: float
The temperature of the system.
mass: float
The mass of the particles.
Returns
-------
momentum_radial: float
The magnitude of the momentum.
"""
momentum_radial = 0
energy = 0.
if temperature > 0.6*mass:
while True:
rand_values = [rd.uniform(0., 1.) for _ in range(3)]
if all(rand_values):
rand_a, rand_b, rand_c = rand_values
momentum_radial = temperature * (rand_a + rand_b + rand_c)
energy = np.sqrt(momentum_radial ** 2. + mass ** 2.)
if rd.uniform(0., 1.) < np.exp((momentum_radial - energy) / temperature):
break
else:
while True:
r0 = rd.uniform(0., 1.)
I1 = mass ** 2.
I2 = 2. * mass * temperature
I3 = 2. * temperature ** 2.
Itot = I1 + I2 + I3
K = 0.0
if r0 < I1 / Itot:
r1 = rd.uniform(0., 1.)
if r1 != 0.:
K = -temperature * np.log(r1)
elif r0 < (I1 + I2) / Itot:
r1, r2 = rd.uniform(0., 1.), rd.uniform(0., 1.)
if r1 != 0. and r2 != 0.:
K = -temperature * np.log(r1 * r2)
else:
r1, r2, r3 = rd.uniform(0., 1.), rd.uniform(0., 1.), rd.uniform(0., 1.)
if r1 != 0. and r2 != 0. and r3 != 0.:
K = -temperature * np.log(r1 * r2 * r3)
energy = K + mass
momentum_radial = np.sqrt((energy + mass) * (energy - mass))
if rd.uniform(0., 1.) < momentum_radial / energy:
break
return momentum_radial
def sample_momenta(self, multiplicity, temperature, mass):
"""
Sample momenta for a given multiplicity, temperature, and mass from
a thermal distribution function.
Parameters
----------
multiplicity: int
The number of particles to sample.
temperature: float
The temperature of the system.
mass: float
The mass of the particles.
Returns
-------
None
"""
p_abs = [self.__thermal_distribution(temperature, mass) for _ in range(multiplicity)]
# compute the directions
azimuths = [self.phi_[p] for p in range(len(p_abs))]
costheta_values = [rd.uniform(-1., 1.) for _ in range(len(p_abs))]
polar_values = np.arccos(costheta_values)
# convert to cartesian
px = [p_abs[p] * np.sin(polar_values[p]) * np.cos(azimuths[p]) for p in range(len(p_abs))]
py = [p_abs[p] * np.sin(polar_values[p]) * np.sin(azimuths[p]) for p in range(len(p_abs))]
pz = [p_abs[p] * np.cos(polar_values[p]) for p in range(len(p_abs))]
self.px_ = px
self.py_ = py
self.pz_ = pz
[docs] def generate_dummy_JETSCAPE_file(self,output_path,number_events,multiplicity,seed):
"""
Generate a dummy JETSCAPE file with random particle momenta.
Parameters
----------
output_path: str
The output file path.
number_events: int
The number of events to generate.
multiplicity: int
The number of particles per event.
seed: int
The random seed for reproducibility.
Returns
-------
None
"""
rd.seed(seed)
temperature = 0.140
mass = 0.138
pdg = 211
status = 27
with open(output_path, "w") as output:
output.write("# JETSCAPE_FINAL_STATE v2 | N pid status E Px Py Pz\n")
for event in range(number_events):
self.sample_angles(multiplicity)
self.sample_momenta(multiplicity, temperature, mass)
output.write(f"# Event {event + 1} weight 1 EPangle 0 N_hadrons {multiplicity}\n")
for particle in range(multiplicity):
energy = np.sqrt(
self.px_[particle] ** 2.
+ self.py_[particle] ** 2.
+ self.pz_[particle] ** 2.
+ mass ** 2.
)
output.write("%d %d %d %g %g %g %g\n"
%(particle,pdg,status,energy,self.px_[particle],
self.py_[particle],self.pz_[particle]))
output.write("# sigmaGen 0.0 sigmaErr 0.0")
[docs] def generate_dummy_OSCAR_file(self,output_path,number_events,multiplicity,seed):
"""
Generate a dummy OSCAR2013 file with random particle momenta.
Parameters
----------
output_path: str
The output file path.
number_events: int
The number of events to generate.
multiplicity: int
The number of particles per event.
seed: int
The random seed for reproducibility.
Returns
-------
None
"""
rd.seed(seed)
temperature = 0.140
mass = 0.138
pdg = 211
status = 27
with open(output_path, "w") as output:
output.write("#!OSCAR2013 particle_lists t x y z mass p0 px py pz pdg ID charge\n")
output.write("# Units: fm fm fm fm GeV GeV GeV GeV GeV none none e\n")
output.write("# SMASH-2.2\n")
for event in range(number_events):
self.sample_angles(multiplicity)
self.sample_momenta(multiplicity, temperature, mass)
output.write(f"# event {event} out {multiplicity}\n")
for particle in range(multiplicity):
energy = np.sqrt(
self.px_[particle] ** 2.
+ self.py_[particle] ** 2.
+ self.pz_[particle] ** 2.
+ mass ** 2.
)
output.write("%g %g %g %g %g %g %g %g %g %d %d %d\n"
%(1,1,1,1,mass, energy,self.px_[particle],
self.py_[particle],self.pz_[particle],pdg, particle, 1))
output.write(f"# event {event} end 0 impact -1.000 scattering_projectile_target no")