Beispiel #1
0
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')
Beispiel #2
0
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')
Beispiel #3
0
        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),
            ]
Beispiel #4
0
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')
Beispiel #5
0
    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
Beispiel #6
0
        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):
Beispiel #7
0
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
Beispiel #8
0
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
Beispiel #9
0
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'])
Beispiel #10
0
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')
Beispiel #11
0
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()
Beispiel #12
0
 def __init__(self, T, **kwargs):
     self._hrgs = [frzout.HRG(t, **kwargs) for t in T]
     self._T4 = T**4 / HBARC**3
Beispiel #13
0
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))