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])