def resonance_mass_distributions(axes):
"""
Verification of mass disributions for several resonance species. Grey
dashed lines are the Breit-Wigner distributions with mass-dependent width,
grey solid lines are the same distributions with momentum integrated out,
and colored lines are histograms of the sampled masses.
"""
with axes() as ax:
T = .15
for ID, name in [
(213, r'$\rho(770)$'),
(2214, r'$\Delta(1232)$'),
(22212, r'$N(1535)$'),
]:
info = frzout.species_dict[ID]
m0 = info['mass']
w0 = info['width']
m_min, m_max = info['mass_range']
sign = -1 if info['boson'] else 1
def bw(m):
w = w0 * np.sqrt((m - m_min) / (m0 - m_min))
return w / ((m - m0)**2 + w * w / 4)
def f(p, m):
return p * p / (np.exp(np.sqrt(p * p + m * m) / T) + sign)
m = np.linspace(m_min, m_max, 200)
ax.plot(m,
bw(m) / integrate.quad(bw, m_min, m_max)[0], **dashed_line)
bwf = np.array([
integrate.quad(lambda p: bw(m_) * f(p, m_), 0, 5)[0]
for m_ in m
]) / integrate.dblquad(lambda m_, p: bw(m_) * f(p, m_), 0, 5,
lambda _: m_min, lambda _: m_max)[0]
ax.plot(m, bwf, color=default_color)
hrg = frzout.HRG(T, species=[ID], res_width=True)
x = np.array([[1, 0, 0, 0]], dtype=float)
sigma = np.array([[1e6 / hrg.density(), 0, 0, 0]])
v = np.zeros((1, 3))
surface = frzout.Surface(x, sigma, v)
parts = frzout.sample(surface, hrg)
m = np.sqrt(np.inner(parts['p']**2, [1, -1, -1, -1]))
ax.hist(m, bins=64, normed=True, histtype='step', label=name)
ax.set_xlim(0, 2)
ax.set_xlabel('Mass [GeV]')
ax.set_ylabel('Probability')
ax.set_yticklabels([])
ax.legend(loc='upper left')
def shear_and_bulk(axes):
"""
Shear tensor with `\pi^{xx} = -\pi^{yy}` and all other components zero, and
bulk pressure fixed at `\Pi = -0.1P_0`. This is a very difficult test case
but the algorithm remains accurate for small to moderate viscous pressure.
"""
hrg = frzout.HRG(.15, res_width=False)
P0 = hrg.pressure()
e0 = hrg.energy_density()
x = np.array([[1., 0, 0, 0]])
sigma = np.array([[1e6 / hrg.density(), 0, 0, 0]])
v = np.zeros((1, 3))
pi_frac = np.linspace(-.5, .5, 11)
Pi_frac = -.1
Tuv = np.array([
sample_Tuv(
frzout.Surface(x,
sigma,
v,
pi=make_pi_dict(xx=i * P0, yy=-i * P0),
Pi=np.array([Pi_frac * P0])), hrg) for i in pi_frac
]).T
P = Tuv.diagonal()[:, 1:].sum(axis=1) / 3
with axes() as ax:
ax.plot(pi_frac, (Tuv[1, 1] - P) / P0, label='$\pi_{xx}$')
ax.plot(pi_frac, pi_frac, **dashed_line)
ax.plot(pi_frac, (Tuv[2, 2] - P) / P0, label='$\pi_{yy}$')
ax.plot(pi_frac, -pi_frac, **dashed_line)
ax.plot(pi_frac, P / P0 - 1, label='Pressure')
ax.axhline(Pi_frac, **dashed_line)
ax.plot(pi_frac, Tuv[0, 0] / e0 - 1, label='Energy density')
ax.axhline(0, **dashed_line)
ax.set_xlim(pi_frac.min(), pi_frac.max())
ax.set_ylim(pi_frac.min(), pi_frac.max())
ax.set_xlabel('$\pi_{xx}/P_0,\ -\pi_{yy}/P_0$')
ax.set_ylabel(
'$\pi_{ij}/P_0,\ \Delta P/P_0,\ \Delta\epsilon/\epsilon_0$')
ax.legend(loc='upper center')
def eos_quantities(T):
hrg = frzout.HRG(T, res_width=False)
parts = frzout.sample(surface, hrg)
E = parts['p'][:, 0]
psq = (parts['p'][:, 1:]**2).sum(axis=1)
T3 = (T / hbarc)**3
T4 = T * T3
return [
(hrg.density() / T3, parts.size / volume / T3),
(hrg.energy_density() / T4, E.sum() / volume / T4),
(3 * hrg.pressure() / T4,
3 * (psq / (3 * E)).sum() / volume / T4),
]
def shear_viscous_corrections(axes):
r"""
Verification that the desired shear tensor `\pi^{\mu\nu}` is reproduced.
This test case checks that nonzero `\pi^{xy}` is reproduced without
changing the equilibrium pressure or energy density. The algorithm starts
to break down at very large shear pressure.
The "shear and bulk" section below has additional checks.
"""
hrg = frzout.HRG(.15, res_width=False)
P0 = hrg.pressure()
e0 = hrg.energy_density()
x = np.array([[1., 0, 0, 0]])
sigma = np.array([[1e6 / hrg.density(), 0, 0, 0]])
v = np.zeros((1, 3))
pi_frac = np.linspace(-.5, .5, 11)
Tuv = np.array([
sample_Tuv(frzout.Surface(x, sigma, v, pi=make_pi_dict(xy=i * P0)),
hrg) for i in pi_frac
]).T
P = Tuv.diagonal()[:, 1:].sum(axis=1) / 3
with axes() as ax:
ax.plot(pi_frac, Tuv[1, 2] / P0, label='$\pi_{xy}$')
ax.plot(pi_frac, pi_frac, **dashed_line)
ax.plot(pi_frac, Tuv[1, 3] / P0, label='$\pi_{xz}$')
ax.plot(pi_frac, P / P0 - 1, label='Pressure')
ax.plot(pi_frac, Tuv[0, 0] / e0 - 1, label='Energy density')
ax.axhline(0, **dashed_line)
ax.set_xlim(pi_frac.min(), pi_frac.max())
ax.set_ylim(pi_frac.min(), pi_frac.max())
ax.set_xlabel('$\pi_{xy}/P_0$')
ax.set_ylabel(
'$\pi_{ij}/P_0,\ \Delta P/P_0,\ \Delta\epsilon/\epsilon_0$')
ax.legend(loc='upper left')
def __init__(self, T, **kwargs):
with warnings.catch_warnings():
warnings.simplefilter('ignore')
self._hrgs = [frzout.HRG(t, **kwargs) for t in T]
self._T4 = T**4 / HBARC**3
resultspce = run_single_event(ic, pce=True)
finalsurface = frzout.Surface(**resultsdict, ymax=2)
finalsurfacepce = frzout.Surface(**resultspce, ymax = 2)
finalresults['initial_entropy'] = ic.sum() * grid_step**2
finalresultspce['initial_entropy'] = ic.sum() * grid_step**2
else:
continue
minsamples, maxsamples = 10, 1000 # reasonable range for nsamples
minparts = 10**5 # min number of particles to sample
nparts = 0 # for tracking total number of sampled particles
npartspce = 0
hrg_kwargs = dict(species='urqmd', res_width=True)
hrg = frzout.HRG(0.150, **hrg_kwargs)
for nsamples in range(1, maxsamples + 1):
parts = frzout.sample(finalsurface, hrg)
if parts.size == 0:
continue
nparts += parts.size
# print('#', parts.size, file=f)
for p in parts:
continue
# print(p['ID'], *p['x'], *p['p'], file=f)
if nparts >= minparts and nsamples >= minsamples:
break
for nsamplespce in range(1, maxsamples + 1):
def run_events(args, results_file, particles_file=None, checkpoint_ic=None):
"""
Run events as determined by user input:
- Read options from `args`, as returned by `parser.parse_args()`.
- Write results to binary file object `results_file`.
- If `checkpoint_ic` is given, run only that IC.
Return True if at least one event completed successfully, otherwise False.
"""
# set the grid step size proportionally to the nucleon width
grid_step = .2#.15*args.nucleon_width
# the "target" grid max: the grid shall be at least as large as the target
grid_max_target = 15
# next two lines set the number of grid cells and actual grid max,
# which will be >= the target (same algorithm as trento)
grid_n = math.ceil(2*grid_max_target/grid_step)
grid_max = .5*grid_n*grid_step
logging.info(
'grid step = %.6f fm, n = %d, max = %.6f fm',
grid_step, grid_n, grid_max
)
def _initial_conditions(nevents=1, initial_file='initial.hdf', avg='off'):
"""
Run trento and yield initial condition arrays.
"""
def average_ic(fname):
with h5py.File(fname, 'a') as f:
densityavg = np.zeros_like(f['event_0/matter_density'].value)
Ncollavg = np.zeros_like(f['event_0/Ncoll_density'].value)
dxy = f['event_0'].attrs['dxy']
Neve = len(f.values())
for eve in f.values():
# step1, center the event
NL = int(eve.attrs['Nx']/2)
density = eve['matter_density'].value
comxy = -np.array(center_of_mass(density))+np.array([NL, NL])
density = shift(density, comxy)
Ncoll = shift(eve['Ncoll_density'].value, comxy)
# step2, rotate the event to align psi2
psi2 = eve.attrs['psi2']
imag_psi2 = psi2*180./np.pi + (90. if psi2<0 else -90.)
densityavg += rotate(density, angle=imag_psi2, reshape=False)
Ncollavg += rotate(Ncoll, angle=imag_psi2, reshape=False)
# step3 take average
densityavg /= Neve
Ncollavg /= Neve
# rewrite the initial.hdf file with average ic
with h5py.File(fname, 'w') as f:
gp = f.create_group('avg_event')
gp.create_dataset('matter_density', data=densityavg)
gp.create_dataset('Ncoll_density', data=Ncollavg)
gp.attrs.create('Nx', densityavg.shape[1])
gp.attrs.create('Ny', densityavg.shape[0])
gp.attrs.create('dxy', dxy)
try:
os.remove(initial_file)
except FileNotFoundError:
pass
if avg == 'on':
logging.info("averaged initial condition mode, could take a while")
run_cmd(
'trento',
'{} {}'.format(args.proj, args.targ),
'--number-events {}'.format(nevents if avg=='off' else 500),
'--grid-step {} --grid-max {}'.format(grid_step, grid_max_target),
'--output', initial_file,
args.trento_args,
)
if avg == 'on':
logging.info("taking average over 500 trento events")
average_ic(initial_file)
### create iterable initial conditon generator
with h5py.File(initial_file, 'r') as f:
for dset in f.values():
ic = np.array(dset['matter_density'])
nb = np.array(dset['Ncoll_density'])
# Write the checkpoint file _before_ starting the event so that
# even if the process is forcefully killed, the state will be
# saved. If / when all events complete, delete the file.
if args.checkpoint is not None:
with open(args.checkpoint, 'wb') as cf:
pickle.dump((args, ic, nb), cf, pickle.HIGHEST_PROTOCOL)
logging.info('wrote checkpoint file %s', args.checkpoint)
yield np.array([ic, nb])
if checkpoint_ic is None:
# if nevents was specified, generate that number of initial conditions
# otherwise generate indefinitely
initial_conditions = (
chain.from_iterable(_initial_conditions() for _ in repeat(None))
if args.nevents is None else
_initial_conditions(args.nevents, avg=args.avg_ic)
)
else:
# just run the checkpointed IC
initial_conditions = [checkpoint_ic[0]]
nbinary_density = [checkpoint_ic[1]]
# create sampler HRG object (to be reused for all events)
hrg_kwargs = dict(species='urqmd', res_width=True)
hrg = frzout.HRG(args.Tswitch, **hrg_kwargs)
# append switching energy density to hydro arguments
eswitch = hrg.energy_density()
hydro_args = [args.hydro_args, 'edec={}'.format(eswitch)]
# arguments for "coarse" hydro pre-runs
# no viscosity, run down to low temperature 110 MeV
hydro_args_coarse = [
'etas_hrg=0 etas_min=0 etas_slope=0 zetas_max=0 zetas_width=0',
'edec={}'.format(frzout.HRG(.110, **hrg_kwargs).energy_density())
]
def save_fs_with_hydro(ic):
# roll ic by index 1 to match hydro
#ic = np.roll(np.roll(ic, shift=-1, axis=0), shift=-1, axis=1)
# use same grid settings as hydro output
with h5py.File('JetData.h5','a') as f:
taufs = f['Event'].attrs['Tau0'][0]
dtau = f['Event'].attrs['dTau'][0]
dxy = f['Event'].attrs['DX'][0]
ls = f['Event'].attrs['XH'][0]
n = 2*ls + 1
coarse = int(dxy/grid_step+.5)
# [tau0, tau0+dtau, tau0+2*dtau, ..., taufs - dtau] + hydro steps...
nsteps = int(taufs/dtau)
tau0 = taufs-dtau*nsteps
if tau0 < 1e-2: # if tau0 too small, skip the first step
tau0 += dtau
nsteps -= 1
taus = np.linspace(tau0, taufs-dtau, nsteps)
# First, rename hydro frames and leave the first few name slots to FS
event_gp = f['Event']
for i in range(len(event_gp.keys()))[::-1]:
old_name = 'Frame_{:04d}'.format(i)
new_name = 'Frame_{:04d}'.format(i+nsteps)
event_gp.move(old_name, new_name)
# Second, overwrite tau0 with FS starting time, and save taufs where
# FS and hydro is separated
event_gp.attrs.create('Tau0', [tau0])
event_gp.attrs.create('TauFS', [taufs])
# Thrid, fill the first few steps with Freestreaming results
for itau, tau in enumerate(taus):
frame = event_gp.create_group('Frame_{:04d}'.format(itau))
fs = freestream.FreeStreamer(ic, grid_max, tau)
for fmt, data, arglist in [
('e', fs.energy_density, [()]),
('V{}', fs.flow_velocity, [(1,), (2,)]),
('Pi{}{}', fs.shear_tensor, [(0,0), (0,1), (0,2),
(1,1), (1,2),
(2,2)] ),
]:
for a in arglist:
X = data(*a).T # to get the correct x-y with vishnew
if fmt == 'V{}': # Convert u1, u2 to v1, v2
X = X/data(0).T
X = X[::coarse, ::coarse]
diff = X.shape[0] - n
start = int(abs(diff)/2)
if diff > 0:
# original grid is larger -> cut out middle square
s = slice(start, start + n)
X = X[s, s]
elif diff < 0:
# original grid is smaller
# -> create new array and place original grid in middle
Xn = np.zeros((n, n))
s = slice(start, start + X.shape[0])
Xn[s, s] = X
X = Xn
if fmt == 'V{}':
Comp = {1:'x', 2:'y'}
frame.create_dataset(fmt.format(Comp[a[0]]), data=X)
if fmt == 'e':
frame.create_dataset(fmt.format(*a), data=X)
frame.create_dataset('P', data=X/3.)
frame.create_dataset('BulkPi', data=X*0.)
prefactor = 1.0/15.62687/5.068**3
frame.create_dataset('Temp', data=(X*prefactor)**0.25)
s = (X + frame['P'].value)/(frame['Temp'].value+1e-14)
frame.create_dataset('s', data=s)
if fmt == 'Pi{}{}':
frame.create_dataset(fmt.format(*a), data=X)
pi33 = -(frame['Pi00'].value + frame['Pi11'].value \
+ frame['Pi22'].value)
frame.create_dataset('Pi33', data=pi33)
pi3Z = np.zeros_like(pi33)
frame.create_dataset('Pi03', data=pi3Z)
frame.create_dataset('Pi13', data=pi3Z)
frame.create_dataset('Pi23', data=pi3Z)
def run_hydro(ic, event_size, coarse=False, dt_ratio=.25):
"""
Run the initial condition contained in FreeStreamer object `fs` through
osu-hydro on a grid with approximate physical size `event_size` [fm].
Return a dict of freeze-out surface data suitable for passing directly
to frzout.Surface.
Initial condition arrays are cropped or padded as necessary.
If `coarse` is an integer > 1, use only every `coarse`th cell from the
initial condition arrays (thus increasing the physical grid step size
by a factor of `coarse`). Ignore the user input `hydro_args` and
instead run ideal hydro down to a low temperature.
`dt_ratio` sets the timestep as a fraction of the spatial step
(dt = dt_ratio * dxy). The SHASTA algorithm requires dt_ratio < 1/2.
"""
# first freestream
fs = freestream.FreeStreamer(ic, grid_max, args.tau_fs)
dxy = grid_step * (coarse or 1)
ls = math.ceil(event_size/dxy) # the osu-hydro "ls" parameter
n = 2*ls + 1 # actual number of grid cells
for fmt, f, arglist in [
('ed', fs.energy_density, [()]),
('u{}', fs.flow_velocity, [(1,), (2,)]),
('pi{}{}', fs.shear_tensor, [(1, 1), (1, 2), (2, 2)]),
]:
for a in arglist:
X = f(*a)
if coarse:
X = X[::coarse, ::coarse]
diff = X.shape[0] - n
start = int(abs(diff)/2)
if diff > 0:
# original grid is larger -> cut out middle square
s = slice(start, start + n)
X = X[s, s]
elif diff < 0:
# original grid is smaller
# -> create new array and place original grid in middle
Xn = np.zeros((n, n))
s = slice(start, start + X.shape[0])
Xn[s, s] = X
X = Xn
X.tofile(fmt.format(*a) + '.dat')
dt = dxy*dt_ratio
run_cmd(
'osu-hydro',
't0={} dt={} dxy={} nls={}'.format(args.tau_fs, dt, dxy, ls),
*(hydro_args_coarse if coarse else hydro_args)
)
surface = np.fromfile('surface.dat', dtype='f8').reshape(-1, 26)
# surface columns:
# 0 1 2 3
# tau x y eta
# 4 5 6 7
# dsigma_t dsigma_x dsigma_y dsigma_z
# 8 9 10
# v_x v_y v_z
# 11 12 13 14
# pitt pitx pity pitz
# 15 16 17
# pixx pixy pixz
# 18 19
# piyy piyz
# 20
# pizz
# 21 22 23 24 25
# Pi T e P muB
if not coarse:
logging.info("Save free streaming history with hydro histroy")
save_fs_with_hydro(ic)
# end event if the surface is empty -- this occurs in ultra-peripheral
# events where the initial condition doesn't exceed Tswitch
if surface.size == 0:
raise StopEvent('empty surface')
# pack surface data into a dict suitable for passing to frzout.Surface
return dict(
x=surface[:, 0:3],
sigma=surface[:, 4:7],
v=surface[:, 8:10],
pi=dict(xx=surface.T[15],xy=surface.T[16], yy=surface.T[18]),
Pi=surface.T[21]
)
results = np.empty((), dtype=result_dtype)
def run_pp_event(ic, nb, event_number):
results.fill(0)
results['initial_entropy'] = ic.sum() * grid_step**2
results['Ncoll'] = nb.sum() * grid_step**2
# Run Pythia+Lido
prefix = os.environ.get('XDG_DATA_HOME')
cmd = "pythia-pp -y {:s}/pythia-pp-setting.txt -i initial.hdf -j {:d} -n {:d}"
run_cmd(cmd.format(prefix, 0 if args.nevents is None else event_number-1, args.NPythiaEvents))
# hadronization
hq = 'c'
prefix = os.environ.get('XDG_DATA_HOME')+"/hvq-hadronization/"
os.environ["ftn20"] = "{}-meson-frzout.dat".format(hq)
os.environ["ftn30"] = prefix+"parameters_{}_hd.dat".format(hq)
os.environ["ftn40"] = prefix+"recomb_{}_tot.dat".format(hq)
os.environ["ftn50"] = prefix+"recomb_{}_BR1.dat".format(hq)
logging.info(os.environ["ftn30"])
subprocess.run("hvq-hadronization", stdin=open("{}-quark-frzout.dat".format(hq)))
cmd = "sed -i -e 's/D/E/g' {}-meson-frzout.dat".format(hq)
subprocess.run(cmd, shell=True)
index, pid, px, py, pz, E, mass, x, y, z, t, T, vx, vy, vz, p0x, p0y, p0z, E0, w, _ \
= np.loadtxt("{}-meson-frzout.dat".format(hq), skiprows=4).T
abs_ID = np.abs(pid)
pT = np.sqrt(px**2+py**2)
pabs = np.sqrt(pT**2 + pz**2)
y = .5*np.log((pabs+pz)/(pabs-pz))
phi = np.arctan2(py, px)
heavy_pid = [411, 421, 413, 423]
is_heavy = np.array([u in heavy_pid for u in abs_ID], dtype=bool)
#============for heavy flavors=======================
for exp in ['ALICE', 'CMS']:
#===========For heavy particles======================
# For charmed hadrons, use info after urqmd
HF_dict = { 'pid': abs_ID[is_heavy],
'pT' : pT[is_heavy],
'y' : y[is_heavy],
'phi': phi[is_heavy],
'w' : w[is_heavy] # normalized to an area units
}
POI = heavy_pid
Yield = JLP.Yield(HF_dict, JEC[exp]['Raa']['pTbins'],
JEC[exp]['Raa']['ybins'], POI)
for s, pid in zip(['D+','D0','D*+', 'D0*'], [411, 421, 413, 423]):
results['dX_dpT_dy_'+exp][s] = Yield[pid][:,0]
#===========For full-pT prediction======================
HF_dict = { 'pid': abs_ID[is_heavy],
'pT' : pT[is_heavy],
'y' : y[is_heavy],
'phi': phi[is_heavy],
'w' : w[is_heavy]
}
POI = heavy_pid
Yield = JLP.Yield(HF_dict, JEC['pred-pT'], [[-1,1]], POI)
for s, pid in zip(['D+','D0','D*+', 'D0*'], [411, 421, 413, 423]):
results['dX_dpT_dy_pred'][s] = Yield[pid][:,0]
def run_single_event(ic, nb, event_number):
"""
Run the initial condition event contained in HDF5 dataset object `ic`
and save observables to `results`.
"""
results.fill(0)
results['initial_entropy'] = ic.sum() * grid_step**2
results['Ncoll'] = nb.sum() * grid_step**2
logging.info("Nb %d", results['Ncoll'])
assert all(n == grid_n for n in ic.shape)
logging.info(
'free streaming initial condition for %.3f fm',
args.tau_fs
)
fs = freestream.FreeStreamer(ic, grid_max, args.tau_fs)
# run coarse event on large grid and determine max radius
rmax = math.sqrt((
run_hydro(ic, event_size=27, coarse=3)['x'][:, 1:3]**2
).sum(axis=1).max())
logging.info('rmax = %.3f fm', rmax)
# now run normal event with size set to the max radius
# and create sampler surface object
surface = frzout.Surface(**run_hydro(ic, event_size=rmax), ymax=2)
logging.info('%d freeze-out cells', len(surface))
# Sampling particle for UrQMD events
logging.info('sampling surface with frzout')
minsamples, maxsamples = 10, 400 # reasonable range for nsamples
minparts = 10**5 # min number of particles to sample
nparts = 0 # for tracking total number of sampled particles
with open('particles_in.dat', 'w') as f:
for nsamples in range(1, maxsamples + 1):
parts = frzout.sample(surface, hrg)
if parts.size == 0:
continue
nparts += parts.size
print('#', parts.size, file=f)
for p in parts:
print(p['ID'], *p['x'], *p['p'], file=f)
if nparts >= minparts and nsamples >= minsamples:
break
results['nsamples'] = nsamples
logging.info('produced %d particles in %d samples', nparts, nsamples)
if nparts == 0:
raise StopEvent('no particles produced')
# ==================Heavy Flavor===========================
# Run Pythia+Lido
prefix = os.environ.get('XDG_DATA_HOME')
run_cmd(
'hydro-couple',
'-y {:s}/pythia-setting.txt'.format(prefix),
'-i ./initial.hdf',
'-j {:d}'.format(0 if args.nevents is None else event_number-1),
'--hydro ./JetData.h5',
'-s {:s}/settings.xml'.format(prefix),
'-t {:s}'.format(args.table_path),
'-n {:d}'.format(args.NPythiaEvents),
args.lido_args,
)
# hadronization
hq = 'c'
prefix = os.environ.get('XDG_DATA_HOME')+"/hvq-hadronization/"
os.environ["ftn20"] = "{}-meson-frzout.dat".format(hq)
os.environ["ftn30"] = prefix+"parameters_{}_hd.dat".format(hq)
os.environ["ftn40"] = prefix+"recomb_{}_tot.dat".format(hq)
os.environ["ftn50"] = prefix+"recomb_{}_BR1.dat".format(hq)
logging.info(os.environ["ftn30"])
subprocess.run("hvq-hadronization", stdin=open("{}-quark-frzout.dat".format(hq)))
# ==================Heavy + Soft --> UrQMD===========================
run_cmd('convert_format {} particles_in.dat c-meson-frzout.dat'.format(nsamples))
run_cmd('run_urqmd urqmd_input.dat particles_out.dat')
# read final particle data
ID, charge, fmass, px, py, pz, y, eta, pT0, y0, w, _ = (
np.array(col, dtype=dtype) for (col, dtype) in
zip(
zip(*read_text_file('particles_out.dat')),
(2*[int] + 10*[float])
)
)
# pT, phi, and id cut
pT = np.sqrt(px**2+py**2)
phi = np.arctan2(py, px)
ET = fmass**2 + pT**2
charged = (charge != 0)
abs_eta = np.fabs(eta)
abs_ID = np.abs(ID)
# It may be redunant to find b-hadron at this stage since UrQMD has
# not included them yet
heavy_pid = [pid for (_, pid) in species.get('heavy')]
is_heavy = np.array([u in heavy_pid for u in abs_ID], dtype=bool)
is_light = np.logical_not(is_heavy)
#============for soft particles======================
results['dNch_deta'] = np.count_nonzero(charged & (abs_eta<.5) & is_light) / nsamples
ET_eta = .6
results['dET_deta'] = ET[abs_eta < ET_eta].sum() / (2*ET_eta) / nsamples
for s, pid in species.get('light'):
cut = (abs_ID == pid) & (abs_eta < 0.5)
N = np.count_nonzero(cut)
results['dN_dy'][s] = N / nsamples
results['mean_pT'][s] = (0. if N == 0 else pT[cut].mean())
pT_alice = pT[charged & (abs_eta < .8) & (.15 < pT) & (pT < 2.)]
results['pT_fluct']['N'] = pT_alice.size
results['pT_fluct']['sum_pT'] = pT_alice.sum()
results['pT_fluct']['sum_pTsq'] = np.inner(pT_alice, pT_alice)
phi_alice = phi[charged & (abs_eta < .8) & (.2 < pT) & (pT < 5.)]
results['Qn_soft']['M'] = phi_alice.size
results['Qn_soft']['Qn'] = [np.exp(1j*n*phi_alice).sum()
for n in range(1, results.dtype['Qn_soft']['Qn'].shape[0] + 1)]
#============for heavy flavors=======================
for exp in ['ALICE', 'CMS']:
#=========Event plane Q-vector from UrQMD events======================
phi_light = phi[charged & is_light \
& (JEC[exp]['vn_ref']['ybins'][0] < eta) \
& (eta < JEC[exp]['vn_ref']['ybins'][1]) \
& (JEC[exp]['vn_ref']['pTbins'][0] < pT) \
& (pT < JEC[exp]['vn_ref']['pTbins'][1])]
results['Qn_ref_'+exp]['M'] = phi_light.shape[0]
results['Qn_ref_'+exp]['Qn'] = np.array([np.exp(1j*n*phi_light).sum()
for n in range(1, 5)])
#===========For heavy particles======================
# For charmed hadrons, use info after urqmd
HF_dict = { 'pid': abs_ID[is_heavy],
'pT' : pT[is_heavy],
'y' : y[is_heavy],
'phi': phi[is_heavy],
'w' : w[is_heavy] # normalized to an area units
}
POI = [pid for (_, pid) in species.get('heavy')]
flow = JLP.Qvector(HF_dict, JEC[exp]['vn_HF']['pTbins'],
JEC[exp]['vn_HF']['ybins'], POI, order=4)
Yield = JLP.Yield(HF_dict, JEC[exp]['Raa']['pTbins'],
JEC[exp]['Raa']['ybins'], POI)
for (s, pid) in species.get('heavy'):
results['dX_dpT_dy_'+exp][s] = Yield[pid][:,0]
results['Qn_poi_'+exp][s]['M'] = flow[pid]['M'][:,0]
results['Qn_poi_'+exp][s]['Qn'] = flow[pid]['Qn'][:,0,:]
# For full pT prediction
#=========Use high precision Q-vector at the end of hydro==============
# oversample to get a high percision event plane at freezeout
ophi_light = np.empty(0)
nloop=0
while ophi_light.size < 10**6 and nloop < 100000:
nloop += 1
oE, opx, opy, opz = frzout.sample(surface, hrg)['p'].T
oM, opT, oy, ophi = JLP.fourvec_to_curvelinear(opx, opy, opz, oE)
ophi = ophi[(-2 < oy) & (oy < 2) & (0.2 < opT) & (opT <5.0)]
ophi_light = np.append(ophi_light, ophi)
results['Qn_ref_pred']['M'] = ophi_light.shape[0]
results['Qn_ref_pred']['Qn'] = np.array([np.exp(1j*n*ophi_light).sum()
for n in range(1, 5)])
del ophi_light
#===========For heavy particles======================
# For charmed hadrons, use info after urqmd
HF_dict = { 'pid': abs_ID[is_heavy],
'pT' : pT[is_heavy],
'y' : y[is_heavy],
'phi': phi[is_heavy],
'w' : w[is_heavy]
}
POI = [pid for (_, pid) in species.get('heavy')]
flow = JLP.Qvector(HF_dict, JEC['pred-pT'], [[-2,2]], POI, order=4)
Yield = JLP.Yield(HF_dict, JEC['pred-pT'], [[-1,1]], POI)
for (s, pid) in species.get('heavy'):
results['dX_dpT_dy_pred'][s] = Yield[pid][:,0]
results['Qn_poi_pred'][s]['M'] = flow[pid]['M'][:,0]
results['Qn_poi_pred'][s]['Qn'] = flow[pid]['Qn'][:,0]
#^^^^^^end of run_single_event(...)^^^^^^^^^^^^^^^^^^
nfail = 0
# run each initial condition event and save results to file
for n, (ic, nb) in enumerate(initial_conditions, start=1):
logging.info('starting event %d', n)
try:
if args.proj == 'p' and args.targ == 'p':
run_pp_event(ic, nb, n)
else:
run_single_event(ic, nb, n)
except StopEvent as e:
if particles_file is not None:
particles_file.create_dataset(
'event_{}'.format(n), shape=(0,), dtype=parts_dtype
)
logging.info('event stopped: %s', e)
except Exception:
logging.exception('event %d failed', n)
nfail += 1
if nfail > 3 and nfail/n > .5:
logging.critical('too many failures, stopping events')
break
logging.warning('continuing to next event')
continue
results_file.write(results.tobytes())
logging.info('event %d completed successfully', n)
# end of events: if running with a checkpoint, delete the file unless this
# was a failed re-run of a checkpoint event
if args.checkpoint is not None:
if checkpoint_ic is not None and nfail > 0:
logging.info(
'checkpoint event failed, keeping file %s',
args.checkpoint
)
else:
os.remove(args.checkpoint)
logging.info('removed checkpoint file %s', args.checkpoint)
return n > nfail
def _realistic_surface_observables():
"""
Compute observables for the "realistic surface" test case.
"""
with open('test_surface.dat', 'rb') as f:
surface_data = np.array(
[l.split() for l in f if not l.startswith(b'#')], dtype=float)
# 0 1 2 3 4 5 6 7
# tau x y dsigma_t dsigma_x dsigma_y v_x v_y
# 8 9 10 11 12 13 14 15
# pitt pitx pity pixx pixy piyy pizz Pi
x, sigma, v, _ = np.hsplit(surface_data, [3, 6, 8])
pixx, pixy, piyy = surface_data.T[11:14]
Pi = surface_data.T[15]
sigma4 = np.zeros((sigma.shape[0], 4))
sigma4[:, :3] = sigma
sigma4 *= x[:, :1]
u_ = np.zeros((v.shape[0], 4))
u_[:, 0] = 1
u_[:, 1:3] = -v
u_ /= np.sqrt(1 - np.square(v).sum(axis=1))[:, np.newaxis]
vx, vy = v.T
pi_uv = np.zeros((pixx.shape[0], 4, 4))
pi_uv[:, 0, 0] = vx * vx * pixx + vy * vy * piyy + 2 * vx * vy * pixy
pi_uv[:, 1, 1] = pixx
pi_uv[:, 2, 2] = piyy
pi_uv[:, 3, 3] = pi_uv[:, 0, 0] - pixx - piyy
pi_uv[:, 0, 1] = pi_uv[:, 1, 0] = -(vx * pixx + vy * pixy)
pi_uv[:, 0, 2] = pi_uv[:, 2, 0] = -(vx * pixy + vy * piyy)
pi_uv[:, 1, 2] = pi_uv[:, 2, 1] = pixy
pT_max = 4
pT_bins = np.linspace(0, pT_max, 41)
pT = (pT_bins[:-1] + pT_bins[1:]) / 2
delta_pT = pT_max / (pT_bins.size - 1)
phi = np.linspace(0, 2 * np.pi, 100, endpoint=False)
eta, eta_weights = special.ps_roots(30)
eta_max = 4
eta *= eta_max
eta_weights *= 2 * eta_max
T = .145
hrg = frzout.HRG(T, res_width=False)
eta_over_tau = hrg.eta_over_tau()
zeta_over_tau = hrg.zeta_over_tau()
cs2 = hrg.cs2()
the_vn = [2, 3, 4]
def calc_obs(ID):
m = frzout.species_dict[ID]['mass']
degen = frzout.species_dict[ID]['degen']
sign = -1 if frzout.species_dict[ID]['boson'] else 1
pT_, phi_, eta_ = np.meshgrid(pT, phi, eta)
mT_ = np.sqrt(m * m + pT_ * pT_)
p = np.array([
mT_ * np.cosh(eta_), pT_ * np.cos(phi_), pT_ * np.sin(phi_),
mT_ * np.sinh(eta_)
]).T
# ignore negative contributions
psigma = np.inner(p, sigma4)
psigma.clip(min=0, out=psigma)
pu = np.inner(p, u_)
with np.errstate(over='ignore'):
f = 1 / (np.exp(pu / T) + sign)
df = f * (1 - sign * f) * (
((pu * pu - m * m) /
(3 * pu) - cs2 * pu) / (zeta_over_tau * T) * Pi +
np.einsum('ijku,ijkv,auv->ijka', p, p, pi_uv) /
(2 * pu * T * eta_over_tau))
f += df
# (phi, pT) distribution
phi_pT_dist = (2 * degen *
np.einsum('i,ijka,ijka->jk', eta_weights, psigma, f) /
(2 * np.pi * hbarc)**3 / phi.size)
pT_dist = phi_pT_dist.sum(axis=1)
# navg, pT dist, qn(pT)
return (2 * np.pi * delta_pT * np.inner(pT, pT_dist), pT_dist, [
np.inner(np.exp(1j * n * phi), phi_pT_dist) / pT_dist
for n in the_vn
])
obs_calc = [calc_obs(i) for i, _ in id_parts]
surface = frzout.Surface(x,
sigma,
v,
pi=dict(xx=pixx, yy=piyy, xy=pixy),
Pi=Pi)
ngroups = 1000
N = 1000 # nsamples per group
nsamples = ngroups * N
# need many samples for diff flow
# too many to store all particles in memory -> accumulate observables
obs_sampled = [
(
np.empty(nsamples, dtype=int), # ID particle counts
np.zeros_like(pT), # pT distribution
np.zeros((len(the_vn), pT.size)), # diff flow
) for _ in id_parts
]
diff_flow_counts = [
np.zeros_like(vn, dtype=int) for (_, _, vn) in obs_sampled
]
from multiprocessing.pool import ThreadPool
for k in range(ngroups):
print(' group', k)
# threading increases performance since sample() releases the GIL
with ThreadPool() as pool:
parts = pool.map(lambda _: frzout.sample(surface, hrg), range(N))
# identified particle counts
for (i, _), (counts, _, _) in zip(id_parts, obs_sampled):
counts[k * N:(k + 1) * N] = [
np.count_nonzero(np.abs(p['ID']) == i) for p in parts
]
# merge all samples
parts = np.concatenate(parts)
abs_ID = np.abs(parts['ID'])
for (i, _), (_, pT_dist, vn_arr), dflow_counts, (_, _, qn_list) in zip(
id_parts, obs_sampled, diff_flow_counts, obs_calc):
parts_ = parts[abs_ID == i]
px, py = parts_['p'].T[1:3]
pT_ = np.sqrt(px * px + py * py)
phi_ = np.arctan2(py, px)
# pT distribution
pT_dist += np.histogram(pT_, bins=pT_bins, weights=1 / pT_)[0]
# differential flow
for n, vn, dfc, qn in zip(the_vn, vn_arr, dflow_counts, qn_list):
cosnphi = [
np.cos(n * phi_[np.fabs(pT_ - p) < .2] - npsi)
for (p, npsi) in zip(pT, np.arctan2(qn.imag, qn.real))
]
vn += [c.sum() for c in cosnphi]
dfc += [c.size for c in cosnphi]
# normalize pT dists and diff flow
for (_, pT_dist, vn), dflow_counts in zip(obs_sampled, diff_flow_counts):
pT_dist /= 2 * np.pi * nsamples * delta_pT
vn /= dflow_counts
return pT, the_vn, obs_calc, obs_sampled
def stress_energy_tensor(axes):
r"""
Verification that the sampling algorithm reproduces the complete
stress-energy tensor `T^{\mu\nu}` from hydrodynamics, including any viscous
corrections.
For each of the following, the flow velocity and viscous pressures are
chosen randomly, particles are sampled, and the effective stress-energy
tensor is computed by summing over the sampled particles. The sampled
tensor is then compared to the expectation from hydro.
In the heatmap cells, the first number is sampled value for the given
component of the tensor and the second number (in parentheses) is the
expected value from hydro. Cells are color-coded, where grey indicates
perfect agreement, red indicates that the sampled value is large, and blue
too small. Some disagreement is expected due to statistical fluctuations
from finite numbers of particles.
Before each heatmap, the randomly chosen velocity and viscous pressures are
listed. The overall magnitude of shear pressure is quantified by the
Lorentz scalar "pirel" = `\sqrt{\pi^{\mu\nu}\pi_{\mu\nu}/(e^2 + 3P_0^2)}`.
"""
hrg = frzout.HRG(.15, res_width=False)
P0 = hrg.pressure()
e0 = hrg.energy_density()
for _ in range(3):
vmag = np.random.rand()
cos_theta = np.random.uniform(-1, 1)
sin_theta = np.sqrt(1 - cos_theta**2)
phi = np.random.uniform(0, 2 * np.pi)
vx = vmag * sin_theta * np.cos(phi)
vy = vmag * sin_theta * np.sin(phi)
vz = vmag * cos_theta
pixx, piyy, pixy, pixz, piyz = np.random.uniform(-.2, .2, 5) * P0
Pi = np.random.uniform(-.3, .3) * P0
surface = frzout.Surface(np.array([[1., 0, 0, 0]]),
np.array([[1e7 / hrg.density(), 0, 0, 0]]),
np.array([[vx, vy, vz]]),
pi={
k[2:]: np.array([v])
for k, v in locals().items()
if k.startswith('pi')
},
Pi=np.array([Pi]))
u = np.array([1, vx, vy, vz]) / np.sqrt(1 - vmag * vmag)
pitt = (vx * vx * pixx + vy * vy * piyy - vz * vz *
(pixx + piyy) + 2 * vx * vy * pixy + 2 * vx * vz * pixz +
2 * vy * vz * piyz) / (1 - vz * vz)
pizz = pitt - pixx - piyy
pitx = vx * pixx + vy * pixy + vz * pixz
pity = vx * pixy + vy * piyy + vz * piyz
pitz = vx * pixz + vy * piyz + vz * pizz
piuv = np.array([
[pitt, pitx, pity, pitz],
[pitx, pixx, pixy, pixz],
[pity, pixy, piyy, piyz],
[pitz, pixz, piyz, pizz],
])
uu = np.outer(u, u)
g = np.array([1, -1, -1, -1], dtype=float)
Delta = np.diag(g) - uu
Tuv_check = e0 * uu - (P0 + Pi) * Delta + piuv
Tuv = u[0] * sample_Tuv(surface, hrg)
Tmag = np.sqrt(e0 * e0 + 3 * P0 * P0)
pimag = np.sqrt(np.einsum('uv,uv,u,v', piuv, piuv, g, g))
diff = (Tuv - Tuv_check) / np.maximum(np.abs(Tuv_check), .1 * Tmag)
tol = .05
fmt = '{:.3f}'
with axes(caption=minus_sign(', '.join([
'v = (' + ', '.join(3 * [fmt]).format(vx, vy, vz) + ')',
'pirel = ' + fmt.format(pimag / Tmag),
'Pi/P0 = ' + fmt.format(Pi / P0),
]))) as ax:
ax.figure.set_size_inches(4.2, 4.2)
ax.figure.set_dpi(100)
ax.imshow(diff, cmap=plt.cm.coolwarm, vmin=-tol, vmax=tol)
for i, j in np.ndindex(*Tuv.shape):
ax.text(i,
j,
minus_sign('\n'.join(
f.format(x[i, j]) for f, x in [
('{:.4f}', Tuv),
('({:.4f})', Tuv_check),
])),
ha='center',
va='center',
fontsize=.75 * font_size)
ax.grid(False)
ax.xaxis.tick_top()
for i in ['x', 'y']:
getattr(ax, 'set_{}ticks'.format(i))(range(4))
getattr(ax, 'set_{}ticklabels'.format(i))(['t', 'x', 'y', 'z'])
def bulk_viscous_corrections(axes):
"""
Effect of bulk viscosity on thermodynamic quantities and momentum
distributions.
The total pressure is the sum of the equilibrium and bulk pressures:
`P = P_0 + \Pi`. In order to satisfy continuity of the stress-energy
tensor, sampled particles must reproduce the total pressure without
changing the equilibrium energy density; this is achieved by parametrically
rescaling overall particle production and momentum depending on the bulk
pressure. This works for negative bulk pressure all the way down to zero
total pressure, but for positive bulk pressure, the necessary momentum
scale factor diverges as the total pressure approaches twice the
equilibrium pressure. The momentum scale is therefore restricted to a
reasonable maximum (3) which effectively limits the positive bulk pressure
to around 70% of the equilibrium pressure, depending on the hadron gas
temperature and composition.
Most quantities are plotted as the relative change to their equilibrium
values vs. the relative bulk pressure `\Pi/P_0`. Colored lines are from
samples and dashed lines are calculated.
"""
T = .15
hrg = frzout.HRG(T, res_width=False)
volume = 1e6 / hrg.density()
x = np.array([[1., 0, 0, 0]])
sigma = np.array([[volume, 0, 0, 0]])
v = np.zeros((1, 3))
def sample_bulk(Pi):
Pi = None if Pi == 0 else np.array([Pi])
surface = frzout.Surface(x, sigma, v, Pi=Pi)
parts = frzout.sample(surface, hrg)
E = parts['p'][:, 0]
psq = (parts['p'][:, 1:]**2).sum(axis=1)
return (parts.size / volume, E.sum() / volume,
(psq / (3 * E)).sum() / volume, np.sqrt(psq).mean())
n0 = hrg.density()
e0 = hrg.energy_density()
P0 = hrg.pressure()
pavg0 = hrg.mean_momentum()
Pi = np.linspace(-P0, P0, 31)
Pi_min, Pi_max = hrg.Pi_lim()
# ensure that the HRG is sampled at precisely the Pi limits
for p in hrg.Pi_lim():
Pi[np.abs(Pi - p).argmin()] = p
Pi_frac = Pi / P0
n, e, P, pavg = np.array([sample_bulk(x) for x in Pi]).T
with axes(
'Pressure and energy density',
'Bulk pressure changes the effective pressure without changing '
'the energy density.') as ax:
ax.plot(Pi_frac, P / P0 - 1, label='Pressure ($P$)')
ax.plot(Pi_frac, Pi_frac.clip(Pi_min / P0, Pi_max / P0), **dashed_line)
ax.plot(Pi_frac, e / e0 - 1, label='Energy density ($\epsilon$)')
ax.axhline(0, **dashed_line)
ax.set_xlim(Pi_frac.min(), Pi_frac.max())
ax.set_ylim(Pi_frac.min(), Pi_frac.max())
ax.set_xlabel('$\Pi/P_0$')
ax.set_ylabel('$\Delta P/P_0$, $\Delta\epsilon/\epsilon_0$')
ax.legend(loc='upper left')
nscale, pscale = np.array([hrg.bulk_scale_factors(p) for p in Pi]).T
with axes(
'Particle density and momentum',
'The changes in particle density and mean momentum necessary '
'to achieve the target pressure and energy density.') as ax:
for y, y0, ycheck, label in [
(n, n0, nscale, 'Density ($n$)'),
(pavg, pavg0, pscale, 'Mean momentum ($p$)'),
]:
ax.plot(Pi_frac, y / y0 - 1, label=label)
ax.plot(Pi_frac, ycheck - 1, **dashed_line)
ax.set_xlim(Pi_frac.min(), Pi_frac.max())
ax.set_xlabel('$\Pi/P_0$')
ax.set_ylabel(r'$\Delta n/n_0$, $\Delta p/p_0$')
ax.legend(loc='upper left')
def f(p, ID):
m, boson, g = (frzout.species_dict[ID][k]
for k in ['mass', 'boson', 'degen'])
s = -1 if boson else 1
return g / (np.exp(np.sqrt(m * m + p * p) / T) + s)
def density(ID, pscale=1):
return (4 * np.pi) / (2 * np.pi * hbarc)**3 * integrate.quad(
lambda p: p * p * f(p / pscale, ID), 0, 10)[0]
ID = 211
n0 = density(ID)
with axes(
'Distribution functions',
'Pion distribution functions `f(p)` for different bulk pressures. '
'Colored histograms are samples; dashed lines are target '
'distributions with rescaled momentum, `f(\lambda p)`.') as ax:
nbins = 50
w = nbins * (2 * np.pi * hbarc)**3 / (2 * volume * 4 * np.pi)
for k, Pi_frac in enumerate([0, -.1, -.3]):
Pi = Pi_frac * P0
parts = frzout.sample(
frzout.Surface(x, sigma, v, Pi=np.array([Pi])), hrg)
psq = (parts[np.abs(parts['ID']) == ID]['p'][:, 1:]**2).sum(axis=1)
pmag = np.sqrt(psq)
offset = 10**(-k)
ax.hist(pmag,
bins=nbins,
weights=w * offset / psq / pmag.ptp(),
histtype='step',
log=True,
label='$\Pi = ' +
('0' if Pi_frac == 0 and k == 0 else
r'{}P_0\ (f \times 10^{{{:d}}})'.format(Pi_frac, -k)) +
'$')
p = np.linspace(0, pmag.max(), 200)
nscale, pscale = hrg.bulk_scale_factors(Pi)
n = density(ID, pscale)
ax.plot(p, n0 / n * nscale * offset * f(p / pscale, ID),
**dashed_line)
ax.set_xlim(0)
ax.set_xlabel('$p\ \mathrm{[GeV]}$')
ax.set_ylabel('$f(p)$')
ax.legend(title='Pions only')
def moving_box(axes):
"""
A single 3D volume element with randomly chosen flow velocity and normal
vector (this randomness means the volume element is not necessarily
realistic for a heavy-ion collision, but it is still numerically valid).
Histograms of sampled (x, y, z) momenta are compared to distributions
computed by numerically integrating the Cooper-Frye function. Negative
contributions are ignored.
"""
T = .15
x = np.array([[1, 0, 0, 0]], dtype=float)
for ID, name in id_parts:
info = frzout.species_dict[ID]
m = info['mass']
g = info['degen']
sign = -1 if info['boson'] else 1
hrg = frzout.HRG(T, species=[ID])
v = np.atleast_2d([np.random.uniform(-i, i) for i in [.5, .5, .7]])
gamma = 1 / np.sqrt(1 - (v * v).sum())
ux, uy, uz = gamma * v.ravel()
volume = 1e6 / hrg.density()
sigma = np.random.uniform(-.5 * volume, 1.5 * volume, (1, 4))
with warnings.catch_warnings():
warnings.filterwarnings('ignore',
'total freeze-out volume is negative')
surface = frzout.Surface(x, sigma, v)
def make_parts():
n = 0
for _ in range(10):
parts = frzout.sample(surface, hrg)
yield parts
n += parts.size
if n > 1e6:
break
parts = list(make_parts())
nsamples = len(parts)
parts = np.concatenate(parts)
psamples = parts['p'].T[1:]
# 3D lattice of momentum points
P = [np.linspace(p.min() - .5, p.max() + .5, 101) for p in psamples]
Px, Py, Pz = np.meshgrid(*P, indexing='ij')
dp = [p.ptp() / (p.size - 1) for p in P]
# evaluate Cooper-Frye function on lattice
E = np.sqrt(m * m + Px * Px + Py * Py + Pz * Pz)
st, sx, sy, sz = sigma.ravel()
dN = ((E * st + Px * sx + Py * sy + Pz * sz) / E / (np.exp(
(E * gamma - Px * ux - Py * uy - Pz * uz) / T) + sign))
dN *= 2 * g / (2 * np.pi * hbarc)**3
# ignore negative contributions
dN.clip(min=0, out=dN)
with axes(name.replace('$', '`') + ' momentum') as ax:
ax.set_yscale('log')
ax.annotate(''.join([
'$\sigma_\mu = (',
', '.join('{:.3f}'.format(i / volume) for i in sigma.flat),
')$\n',
'$v = (',
', '.join('{:.3f}'.format(i) for i in v.flat),
')$',
]), (.03, .96),
xycoords='axes fraction',
ha='left',
va='top')
nbins = 50
for i, (p, c) in enumerate(zip(psamples, ['x', 'y', 'z'])):
ax.hist(p,
bins=nbins,
weights=np.full_like(p, nbins / p.ptp() / nsamples),
histtype='step',
label='$p_{}$'.format(c))
j, k = set(range(3)) - {i}
# evaluate dN/dp_i by integrating out axes (j, k)
ax.plot(P[i],
dp[j] * dp[k] * dN.sum(axis=(j, k)),
color=default_color)
ax.set_xlabel('$p\ [\mathrm{GeV}]$')
ax.set_ylabel('$dN/dp\ [\mathrm{GeV}^{-1}]$')
ax.yaxis.get_major_locator().base(100)
ax.legend()
def __init__(self, T, **kwargs):
self._hrgs = [frzout.HRG(t, **kwargs) for t in T]
self._T4 = T**4 / HBARC**3
def main():
collision_sys = 'PbPb5020'
spectraFile = '%s/spectra/LHC5020-AA2ccbar.dat' % share
# ==== parse the config file ============================================
if len(sys.argv) == 3:
config = parseConfig(sys.argv[1])
jobID = sys.argv[2]
else:
config = {}
jobID = 0
# ====== set up grid size variables ======================================
grid_step = 0.1
grid_max = 15.05
dtau = 0.25 * grid_step
Nhalf = int(grid_max / grid_step)
tau_fs = float(config.get('tau_fs'))
xi_fs = float(config.get('xi_fs'))
nevents = int(config.get('nevents'))
# ========== initial condition ============================================
proj = collision_sys[:2]
targ = collision_sys[2:4]
run_cmd('trento {} {}'.format(proj, targ), str(nevents),
'--grid-step {} --grid-max {}'.format(grid_step, grid_max),
'--output {}'.format('initial.hdf5'),
config.get('trento_args', ''))
run_qhat(config.get('qhat_args'))
# set up sampler HRG object
Tswitch = float(config.get('Tswitch'))
hrg = frzout.HRG(Tswitch, species='urqmd', res_width=True)
eswitch = hrg.energy_density()
finitial = h5py.File('initial.hdf5', 'r')
for (ievent, dset) in enumerate(finitial.values()):
resultFile = 'result_{}-{}.hdf5'.format(jobID, ievent)
fresult = h5py.File(resultFile, 'w')
print('# event: ', ievent)
ic = [dset['matter_density'].value, dset['Ncoll_density'].value]
event_gp = fresult.create_group('initial')
event_gp.attrs.create('initial_entropy', grid_step**2 * ic[0].sum())
event_gp.attrs.create('N_coll', grid_step**2 * ic[1].sum())
for (k, v) in list(finitial['event_{}'.format(ievent)].attrs.items()):
event_gp.attrs.create(k, v)
# =============== Freestreaming ===========================================
save_fs_history(ic[0],
event_size=grid_max,
grid_step=grid_step,
tau_fs=tau_fs,
xi=xi_fs,
steps=5,
grid_max=grid_max,
coarse=2)
fs = freestream.FreeStreamer(ic[0], grid_max, tau_fs)
e = fs.energy_density()
e_above = e[e > eswitch].sum()
event_gp.attrs.create('multi_factor',
e.sum() / e_above if e_above > 0 else 1)
e.tofile('ed.dat')
# calculate the participant plane angle
participant_plane_angle(e, int(grid_max))
for i in [1, 2]:
fs.flow_velocity(i).tofile('u{}.dat'.format(i))
for ij in [(1, 1), (1, 2), (2, 2)]:
fs.shear_tensor(*ij).tofile('pi{}{}.dat'.format(*ij))
# ============== vishnew hydro ===========================================
run_cmd(
'vishnew initialuread=1 iein=0',
't0={} dt={} dxy={} nls={}'.format(tau_fs, dtau, grid_step, Nhalf),
config.get('hydro_args', ''))
# ============= frzout sampler =========================================
surface_data = np.fromfile('surface.dat', dtype='f8').reshape(-1, 16)
if surface_data.size == 0:
print("empty event")
continue
print('surface_data.size: ', surface_data.size)
surface = frzout.Surface(**dict(
zip(['x', 'sigma', 'v'], np.hsplit(surface_data, [3, 6, 8])),
pi=dict(zip(['xx', 'xy', 'yy'], surface_data.T[11:14])),
Pi=surface_data.T[15]),
ymax=3.)
minsamples, maxsamples = 10, 100
minparts = 30000
nparts = 0 # for tracking total number of sampeld particles
# sample soft particles and write to file
with open('particle_in.dat', 'w') as f:
nsamples = 0
while nsamples < maxsamples + 1:
parts = frzout.sample(surface, hrg)
if parts.size == 0:
continue
else:
nsamples += 1
nparts += parts.size
print("#", parts.size, file=f)
for p in parts:
print(p['ID'],
*itertools.chain(p['x'], p['p']),
file=f)
if nparts >= minparts and nsamples >= minsamples:
break
event_gp.attrs.create('nsamples', nsamples, dtype=np.int)
# =============== HQ initial position sampling ===========================
initial_TAA = ic[1]
np.savetxt('initial_Ncoll_density.dat', initial_TAA)
HQ_sample_conf = {'IC_file': 'initial_Ncoll_density.dat',\
'XY_file': 'initial_HQ.dat', \
'IC_Nx_max': initial_TAA.shape[0], \
'IC_Ny_max': initial_TAA.shape[1], \
'IC_dx': grid_step, \
'IC_dy': grid_step, \
'IC_tau0': 0, \
'N_sample': 60000, \
'N_scale': 0.05, \
'scale_flag': 0}
ftmp = open('HQ_sample.conf', 'w')
for (key, value) in zip(HQ_sample_conf.keys(),
HQ_sample_conf.values()):
inputline = ' = '.join([str(key), str(value)]) + '\n'
ftmp.write(inputline)
ftmp.close()
run_cmd('HQ_sample HQ_sample.conf')
# ================ HQ evolution (pre-equilibirum stages) =================
os.environ['ftn00'] = 'FreeStream.h5'
os.environ['ftn10'] = '%s/dNg_over_dt_cD6.dat' % share
print(os.environ['ftn10'])
os.environ['ftn20'] = 'HQ_AAcY_preQ.dat'
os.environ['ftn30'] = 'initial_HQ.dat'
run_cmd('diffusion hq_input=3.0 initt={}'.format(tau_fs * xi_fs),
config.get('diffusion_args', ''))
# ================ HQ evolution (in medium evolution) ====================
os.environ['ftn00'] = 'JetData.h5'
os.environ['ftn10'] = '%s/dNg_over_dt_cD6.dat' % share
os.environ['ftn20'] = 'HQ_AAcY.dat'
os.environ['ftn30'] = 'HQ_AAcY_preQ.dat'
run_cmd('diffusion hq_input=4.0 initt={}'.format(tau_fs),
config.get('diffusion_args', ''))
# ============== Heavy quark hardonization ==============================
os.environ['ftn20'] = 'Dmeson_AAcY.dat'
child1 = 'cat HQ_AAcY.dat'
p1 = subprocess.Popen(child1.split(), stdout=subprocess.PIPE)
p2 = subprocess.Popen('fragPLUSrecomb', stdin=p1.stdout)
p1.stdout.close()
output = p2.communicate()[0]
# ============ Heavy + soft UrQMD =================================
run_cmd(
'afterburner {} urqmd_final.dat particle_in.dat Dmeson_AAcY.dat'.
format(nsamples))
# =========== processing data ====================================
calculate_beforeUrQMD(spectraFile, 'Dmeson_AAcY.dat', resultFile,
'beforeUrQMD/Dmeson', 1.0, 'a')
calculate_beforeUrQMD(spectraFile, 'HQ_AAcY.dat', resultFile,
'beforeUrQMD/HQ', 1.0, 'a')
calculate_beforeUrQMD(spectraFile, 'HQ_AAcY_preQ.dat', resultFile,
'beforeUrQMD/HQ_preQ', 1.0, 'a')
if nsamples != 0:
calculate_afterUrQMD(spectraFile, 'urqmd_final.dat', resultFile,
'afterUrQMD/Dmeson', 1.0, 'a')
shutil.move('urqmd_final.dat',
'urqmd_final{}-{}.dat'.format(jobID, ievent))
shutil.move('Dmeson_AAcY.dat',
'Dmeson_AAcY{}-{}.dat'.format(jobID, ievent))
shutil.move('HQ_AAcY.dat', 'HQ_AAcY{}-{}.dat'.format(jobID, ievent))
shutil.move('HQ_AAcY_preQ.dat',
'HQ_AAcY_preQ{}-{}.dat'.format(jobID, ievent))
#=== after everything, save initial profile (depends on how large the size if, I may choose to forward this step)
shutil.move('initial.hdf5', 'initial_{}.hdf5'.format(jobID))