def run_single_event(ic, pce=False):
fs = freestream.FreeStreamer(ic, grid_max, 0.5)
surface = run_hydro(fs, event_size=27, coarse=3)
resdict = dict(zip(['x', 'sigma', 'v'], np.hsplit(surface, [3,6,8])))
rmax = math.sqrt((resdict['x'][:,1:3]**2).sum(axis=1).max())
if pce:
results = run_hydro_pce(fs, event_size=rmax)
else:
results = run_hydro(fs, event_size=rmax)
return results
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]
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]
)
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))
def __call__(self, n):
t = n * self.dt
print('frame {:d} / {:d}, t = {:1.2f} fm'.format(n, self.nframes, t))
fs = freestream.FreeStreamer(self.initial, self.grid_max, t)
return t, fs.Tuv(0, 0)
def save_fs_history(ic,
event_size,
grid_step,
tau_fs,
xi,
grid_max,
steps=5,
coarse=False):
f = h5py.File('FreeStream.h5', 'w')
dxy = grid_step * (coarse or 1)
ls = math.ceil(event_size / dxy)
n = 2 * ls + 1
NX, NY = ic.shape
# roll ic by index 1 to match hydro
ix = np.roll(np.roll(ic, shift=-1, axis=0), shift=-1, axis=1)
tau0 = tau_fs * xi
taus = np.linspace(tau0, tau_fs, steps)
dtau = taus[1] - taus[0]
gp = f.create_group('Event')
gp.attrs.create('XL', [-ls])
gp.attrs.create('XH', [ls])
gp.attrs.create('YL', [-ls])
gp.attrs.create('YH', [ls])
gp.attrs.create('Tau0', [tau0])
gp.attrs.create('dTau', [dtau])
gp.attrs.create('DX', [dxy])
gp.attrs.create('DY', [dxy])
gp.attrs.create('NTau', [steps])
gp.attrs.create('OutputViscousFlag', [1])
for itau, tau in enumerate(taus):
print(tau)
frame = 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?? (need to check this)
if fmt == 'V{}':
X = X / data(0).T
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
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)
f.close()
def test_gaussian():
"""
check symmetric gaussian against analytic solution
"""
# sample random thermalization time and Gaussian width
t0 = np.random.uniform(0.01, 1.5)
sigma = np.random.uniform(0.4, 1.0)
sigma_sq = sigma*sigma
xymax = 5*sigma + 2*t0
nsteps = 2*int(xymax/0.1) + 1
s = np.s_[-xymax:xymax:nsteps*1j]
Y, X = np.mgrid[s, s]
grid_max = xymax/(1 - 1/nsteps)
for var in ['t0', 'sigma', 'xymax', 'nsteps']:
print(var, '=', eval(var))
initial = np.exp(-(X*X + Y*Y)/(2*sigma_sq)) / (2*np.pi*sigma_sq)
fs = freestream.FreeStreamer(initial, grid_max, t0)
initial_sum = initial.sum()
final_sum = fs.Tuv(0, 0).sum() * t0
assert abs(initial_sum/final_sum - 1) < 1e-6, \
'particle number is not conserved: {} != {}'.format(initial_sum,
final_sum)
# compare all quantities to analytic solution along the positive x-axis
# skip the outermost cells where everything is tiny
pos_x = np.s_[int(nsteps/2), int(nsteps/2)+1:int(0.95*nsteps)]
x = X[pos_x]
# T^uv can be computed analytically on the x-axis starting from a symmetric
# Gaussian initial state. After appropriate change of the variables the
# integrals become Bessel functions.
# This prefactor multiplies the entire matrix.
prefactor = np.exp(-(x*x + t0*t0)/(2*sigma_sq)) / (2*np.pi*sigma_sq*t0)
# Dimensionless variable that appears often.
w = x*t0/sigma_sq
# Bessel functions I_0 and I_1.
i0 = special.i0(w)
i1 = special.i1(w)
# Compare to exact results for T^uv.
for (u, v, Tuv_exact) in [
(0, 0, i0),
(0, 1, i1),
(0, 2, 0),
(1, 1, i0 - i1/w),
(1, 2, 0),
(2, 2, i1/w),
]:
assert_allclose(fs.Tuv(u, v)[pos_x], prefactor*Tuv_exact,
'T{}{}'.format(u, v))
# The Landau matching eigenvalue problem can also be solved analytically.
# This discriminant "d" from the characteristic equation appears frequently
# in the remaining expressions.
d = np.sqrt(i1*i1 - 4*i0*i1*w + 4*(i0-i1)*(i0+i1)*w*w)
# Verify energy density (eigenvalue).
energy_density = prefactor/(2*w) * (i1 + d)
assert_allclose(fs.energy_density()[pos_x], energy_density,
'energy density')
# Verify flow velocity (eigenvector).
u0, u1, u2 = fs.flow_velocity().T
assert np.allclose(u0*u0 - u1*u1 - u2*u2, 1), \
'flow velocities are not normalized'
v0 = i0/i1 - 1/(2*w) + d/(2*i1*w)
v1 = 1
v2 = 0
v_norm = np.sqrt(v0*v0 - v1*v1 - v2*v2)
for i, v in enumerate([v0, v1, v2]):
assert_allclose(fs.flow_velocity(i)[pos_x], v/v_norm,
'flow velocity u{}'.format(i))
# The shear tensor pi^uv can also be computed analytically, although the
# expressions are somewhat tedious...
for (u, v, piuv_exact) in [
(0, 0, (i1*i1*(1 + 4*w*w) + i1*d + 2*i0*w*(d - 2*i0*w))/(6*w*d)),
(0, 1, i1/3 * (1 - 2*i1/d)),
(0, 2, 0),
(1, 1, (i1*i1*(3 - 4*w*w) + 2*i0*w*(2*i0*w + d) -
i1*(8*i0*w + 3*d))/(6*w*d)),
(1, 2, 0),
(2, 2, (5*i1 - d)/(6*w)),
]:
assert_allclose(fs.shear_tensor(u, v)[pos_x], prefactor*piuv_exact,
'shear tensor pi{}{}'.format(u, v))
_check_shear_orthogonal(fs)
# Bulk pressure is zero for ideal eos...
assert_allclose(fs.bulk_pressure(), 0, 'ideal bulk pressure',
rtol=1e-5, atol=1e-10)
# ...and nonzero for any other eos.
assert_allclose(
fs.bulk_pressure(eos=lambda e: e/6)[pos_x], energy_density/6,
'nonideal bulk pressure', rtol=1e-4, atol=1e-7
)
def main():
parser = argparse.ArgumentParser(
description='plot freestream test cases')
parser.add_argument('plot', choices=['gaussian1', 'gaussian2', 'random'],
help='test case to plot')
parser.add_argument('output', help='plot output file')
args = parser.parse_args()
grid_max = 5.0
nsteps = 100
xymax = grid_max*(1 - 1/nsteps)
s = np.s_[-xymax:xymax:nsteps*1j]
Y, X = np.mgrid[s, s]
if args.plot == 'gaussian1':
initial = np.exp(-(X*X + Y*Y)/(2*0.5**2))
elif args.plot == 'gaussian2':
initial = np.exp(-(np.square(X - 0.5) + 3*np.square(Y - 1)))
elif args.plot == 'random':
Y, X = np.mgrid[s, s]
initial = np.zeros_like(X)
sigmasq = .4**2
# truncate gaussians at several widths (mimics typical IC models)
rsqmax = 5**2 * sigmasq
for x0, y0 in np.random.standard_normal((25, 2)):
rsq = (X - x0)**2 + (Y - y0)**2
cut = rsq < rsqmax
initial[cut] += np.exp(-.5*rsq[cut]/sigmasq)
fs = freestream.FreeStreamer(initial, grid_max, 1.0)
plt.rcdefaults()
def pcolorfast(arr, ax=None, title=None, vrange=None):
if ax is None:
ax = plt.gca()
arrmax = np.abs(arr).max()
try:
vmin, vmax = vrange
except TypeError:
vmax = arrmax if vrange is None else vrange
vmin = -vmax
else:
vmin = 2*vmin - vmax
ax.set_aspect('equal')
ax.pcolorfast((-grid_max, grid_max), (-grid_max, grid_max), arr,
vmin=vmin, vmax=vmax, cmap=plt.cm.RdBu_r)
ax.text(-0.9*grid_max, 0.9*grid_max, 'max = {:g}'.format(arrmax),
ha='left', va='top')
ax.set_xlim(-grid_max, grid_max)
ax.set_ylim(-grid_max, grid_max)
if title is not None:
ax.set_title(title)
with PdfPages(args.output) as pdf:
def finish():
plt.tight_layout()
pdf.savefig()
plt.close()
gamma_max = np.percentile(fs.flow_velocity(0), 90)
for arr, title, vrange in [
(initial, 'initial state', None),
(fs.energy_density(), r'energy density', None),
(fs.flow_velocity(0), r'$u^0$', (1, gamma_max)),
(fs.flow_velocity(1), r'$u^1$', gamma_max),
(fs.flow_velocity(2), r'$u^2$', gamma_max),
(fs.bulk_pressure(), r'$\Pi$', 1e-15),
]:
plt.figure(figsize=(6, 6))
pcolorfast(arr, title=title, vrange=vrange)
plt.xlabel(r'$x$ [fm]')
plt.ylabel(r'$y$ [fm]')
finish()
for (tensor, name) in [(fs.Tuv, 'T'), (fs.shear_tensor, r'\pi')]:
fig, axes = plt.subplots(nrows=3, ncols=3,
sharex='col', sharey='row',
figsize=(12, 12))
for (u, v), ax in np.ndenumerate(axes):
if u < v:
ax.set_axis_off()
else:
pcolorfast(tensor(u, v), ax,
r'${}^{{{}{}}}$'.format(name, u, v))
if ax.is_last_row():
ax.set_xlabel(r'$x$ [fm]')
if ax.is_first_col():
ax.set_ylabel(r'$y$ [fm]')
finish()
def test_random():
"""
check random initial condition against non-interpolated solution
"""
for bad_shape in [10, (10, 12), (10, 10, 2)]:
with assert_raises(ValueError):
freestream.FreeStreamer(np.empty(bad_shape), 10, 1)
# Test the algorithm on a more interesting initial state that cannot be
# solved analytically.
# Construct a random initial state by sampling normally-distributed
# positions for Gaussian blobs.
t0 = np.random.uniform(0.01, 1.5)
sigma = np.random.uniform(0.4, 0.8)
xy0 = np.random.standard_normal(50).reshape(-1, 2)
# Define function that evaluates the initial density at (x, y) points.
def f(x, y):
z = np.zeros_like(x)
for (x0, y0) in xy0:
z += np.exp(-(np.square(x - x0) + np.square(y - y0))/(2*sigma**2))
z /= 2*np.pi*sigma**2
return z
# Discretize the function onto a grid.
xymax = np.max(xy0) + 5*sigma + 2*t0
nsteps = int(2*xymax/0.1) + 2
grid_max = xymax/(1 - 1/nsteps)
s = np.s_[-xymax:xymax:nsteps*1j]
Y, X = np.mgrid[s, s]
for var in ['t0', 'sigma', 'xymax', 'grid_max', 'nsteps']:
print(var, '=', eval(var))
initial = f(X, Y)
fs = freestream.FreeStreamer(initial, grid_max, t0)
initial_sum = initial.sum()
final_sum = fs.Tuv(0, 0).sum() * t0
assert abs(initial_sum/final_sum - 1) < 1e-6, \
'particle number is not conserved: {} != {}'.format(initial_sum,
final_sum)
# Pick some random grid points to check near the middle-ish of the grid.
check_indices = np.random.randint(0.25*nsteps, 0.75*nsteps, (10, 2))
check_Tuv = fs.Tuv()[check_indices.T[1], check_indices.T[0]]
check_xy = (check_indices + 0.5)*2*grid_max/nsteps - grid_max
print('check indices and points:')
for (ix, iy), (x, y) in zip(check_indices, check_xy):
print('{: 5d}{: 4d}{: 7.2f}{: 6.2f}'.format(ix, iy, x, y))
# Check the components of T^uv against adaptive quadrature integration of
# the actual continuous function.
for u, v, w in [
(0, 0, lambda phi: 1),
(0, 1, np.cos),
(0, 2, np.sin),
(1, 1, lambda phi: np.square(np.cos(phi))),
(1, 2, lambda phi: np.cos(phi)*np.sin(phi)),
(2, 2, lambda phi: np.square(np.sin(phi))),
]:
approx = check_Tuv[:, u, v] * t0
exact = [
integrate.quad(
lambda phi: f(x0 - t0*np.cos(phi), y0 - t0*np.sin(phi))*w(phi),
0, 2*np.pi
)[0]/(2*np.pi) for (x0, y0) in check_xy
]
assert_allclose(approx, exact, 'T{}{}'.format(u, v))
# check basic class functionality and properties
# check energy-momentum tensor
T00 = fs.Tuv(0, 0)
assert all([
T00.base is fs.Tuv(),
T00.shape == fs.Tuv().shape[:2],
np.all(T00 == fs.Tuv()[..., 0, 0]),
]), 'T00 is not a view of Tuv'
assert_raises(ValueError, fs.Tuv, 0)
# check flow velocity
u0, u1, u2 = fs.flow_velocity().transpose(2, 0, 1)
assert all([
u1.base is fs.flow_velocity(),
u1.shape == fs.flow_velocity().shape[:2],
np.all(u1 == fs.flow_velocity()[..., 1]),
]), 'u1 is not a view of flow_velocity'
assert np.allclose(u0*u0 - u1*u1 - u2*u2, 1), \
'flow velocities are not normalized'
# check shear tensor
pi11 = fs.shear_tensor(1, 1)
assert all([
pi11.base is fs.shear_tensor(),
pi11.shape == fs.shear_tensor().shape[:2],
np.all(pi11 == fs.shear_tensor()[..., 1, 1]),
]), 'pi11 is not a view of shear_tensor'
_check_shear_orthogonal(fs)
assert_raises(ValueError, fs.shear_tensor, 2)
# check bulk pressure
assert_allclose(fs.bulk_pressure(), 0, 'ideal bulk pressure',
rtol=1e-5, atol=1e-10)
assert_allclose(
fs.bulk_pressure(eos=lambda e: e/6), fs.energy_density()/6,
'nonideal bulk pressure', rtol=1e-5, atol=1e-10
)
# check definition of Tuv
e = fs.energy_density()[..., np.newaxis, np.newaxis]
uu = np.einsum('...i,...j', fs.flow_velocity(), fs.flow_velocity())
Delta = np.diag([1., -1., -1.]) - uu
Peff = e/3 + fs.bulk_pressure()[..., np.newaxis, np.newaxis]
Tuv_calc = e*uu - Peff*Delta + fs.shear_tensor()
assert np.allclose(fs.Tuv(), Tuv_calc), \
'Tuv does not match the sum of its parts'
def run_hydro(initial_density,
args,
ic_grid_max=15.0,
hydro_grid_max=20.0,
dxy=0.1):
# first freestream
fs = freestream.FreeStreamer(initial_density, ic_grid_max, args.tau_fs)
ls = math.ceil(hydro_grid_max / 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)
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 * 0.25
run_cmd('osu-hydro',
't0={} dt={} dxy={} nls={}'.format(args.tau_fs, dt, dxy,
ls), args.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
logging.info("Save free streaming history with hydro histroy")
save_fs_with_hydro(initial_density, ic_grid_max)
# 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])
