Machine Learning Approach to Analyze the Heavy Quark Diffusion Coefficient in Relativistic Heavy Ion Collisions

Abstract

The diffusion coefficient of heavy quarks in a deconfined medium is examined in this research using a deep convolutional neural network (CNN) that is trained with data from relativistic heavy ion collisions involving heavy flavor hadrons. The CNN is trained using observables such as the nuclear modification factor $R_{AA}$ and the elliptic flow $v_2$ of non-prompt $J/\psi$ from the B-hadron decay in different centralities, where B meson evolutions are calculated using the Langevin equation and the instantaneous coalescence model. The CNN outputs the parameters, thereby characterizing the temperature and momentum dependence of the heavy quark diffusion coefficient. By inputting the experimental data of the non-prompt $J/\psi$$(R_{AA},v_2)$ from various collision centralities into multiple channels of a well-trained network, we derive the values of the diffusion coefficient parameters. Additionally, we evaluate the uncertainty in determining the diffusion coefficient by taking into account the uncertainties present in the experimental data $(R_{AA},v_2)$, which serve as inputs to the deep neural network.

1. Introduction

In recent years, there has been rapid development in deep learning methods, which are increasingly being widely applied in industry and scientific research. In particular, deep learning methods are being used to handle high-dimensional data and to uncover patterns, such as in image recognition. In the realm of scientific research, deep learning has already found applications in many aspects of physics [1,2,3]. In the theoretical research of high-energy nuclear physics, an increasing number of studies are utilizing deep learning methods to analyze the computational data from theoretical models, such as the study of the equation of the state of QGP matter [4], the dynamical evolutions of QGP [5] and identifications of the spinodal clumping in high-energy nuclear collisions [6].

In the Relativistic Heavy-Ion Collider (RHIC) and the Large Hadron Collider (LHC) [7,8,9], a new kind of deconfined state, called the Quark-Gluon Plasma (QGP), was predicted and produced. Properties such as the initial energy density and coupling strength of QGP have been typically studied through the final-state light hadron spectra [10,11,12,13] or through the distribution of heavy flavor particles [14,15,16,17,18,19]. For heavy flavor particles, their distribution is initially influenced by the cold nuclear effect, and then they will interact with the QGP medium, thus resulting in energy loss [20,21,22,23]. Numerous models have been established to account for the various effects mentioned above [24,25,26,27,28,29], whereby the aim is to ultimately predict the nuclear modification factor and the anisotropic flows of open heavy flavor hadrons. These models have been used to investigate the energy loss mechanisms of heavy quarks in the QGP from different perspectives. Relevant nuclear modification factors and the collective flows of D mesons or B mesons have also been experimentally measured by STAR [30], ALICE [31,32,33] and CMS [34,35] collaborations, which are closely related to the energy loss process of heavy quarks. As it is not straightforward to explain the experimental observables of $R_{AA}$ and $v_2\left(p_T\right)$ at the same time when using a simple value for the diffusion coefficient $D_{s}2\pi T$, a more realistic expression with temperature and momentum dependence is needed. Due to the complex process involving the heavy quark diffusion and hydrodynamic evolution, it is necessary to employ deep neural networks to analyze the relationship between the diffusion coefficient and experimental observables. CNN has been proven to be suitable for analyzing high-dimensional datasets and for quantifying the value of a diffusion coefficient when considering multiple hot and cold nuclear matter effects.

In previous studies, Bayesian statistical analysis has been employed to analyze the experimental data of soft particles [36,37,38,39], whereby the diffusion coefficient is quantitatively extracted with experimental observables of the charmed hadrons in heavy-ion collisions [40]. Although there is abundant experimental data on D mesons, our main focus is on studying the evolution of B mesons. This is because the diffusion coefficient is directly related to the drag term in the Langevin equation, which arises from elastic scattering processes rather than medium-induced gluon radiation. Therefore, as the mass of a heavy quark increases, the contribution from elastic scattering processes becomes more significant in the energy loss process, while the contribution from gluon radiation is relatively weaker. In this work, we treat the $R_{AA}$ and $v_2$ of non-prompt $J/\psi$ from the B meson decay as inputs of the CNN, and the parameters in the diffusion coefficient are treated as outputs of the network. After training the neural network with supervision, the inputs of the neural network are selected from values within the error bars of the experimental data points in order to generate the corresponding diffusion coefficient (albeit with some uncertainty).

This paper is organized as follows: in Section 2, we introduce the Langevin plus instantaneous coalescence model (LICM) to generate datasets of heavy flavor evolutions that had different values in terms of their parameters (which were then used as training datasets for the CNN model). In Section 3, the values of the shadowing factor, temperature and momentum dependence of the diffusion coefficient are quantitatively extracted based on the experimental data of the non-prompt $J/\psi$ from B meson decay. A final summary is given in Section 4.

2. Frameworks

2.1. Generating Datasets with LICM

Bottom quarks are produced in the hard scatterings of nucleon partons. We adopted the fixed-order plus next-to-leading log formula (FONLL) [41,42], and the NNPDF30NLO PDF set [43] was used to calculated the initial momentum distribution of the bottom quarks in nucleon–nucleon collisions. Due to the fact that Pb–Pb collisions can be regarded as a superposition of nucleon–nucleon collisions, which are accompanied by cold nuclear matter effects, the initial momentum distribution of the bottom quarks in Pb–Pb collisions can be considered as the momentum distribution in the pp collisions that are multiplied by a shadowing factor. The nuclear shadowing factor was calculated with the EPS09 NLO package [44] in 5.02 TeV Pb-Pb collisions. The production of the bottom quarks primarily arose from binary collision processes; hence, the spatial distribution of the bottom quarks in nuclear collisions was proportional to the distribution of binary collisions $d{N}_{b\overline{b}}/d{\mathbf{x}}_T\propto T_A({\mathbf{x}}_T-\mathbf{b}/2)T_B({\mathbf{x}}_T+\mathbf{b}/2)$ [45]. Here, $T_{A\left(B\right)}=\int dz{\rho }_{A\left(B\right)}({\mathbf{x}}_T,z)$ represents the thickness functions of the two nuclei. The nucleon distribution $\rho ({\mathbf{x}}_T,z)$ in the nucleus follows a Woods–Saxon distribution.

After the production of bottom quarks, they propagate within the high-temperature QGP medium and are accompanied by energy loss. The energy loss of bottom quarks in the QGP is primarily attributed to the scattering processes between bottom quarks and thermal partons, as well as medium-induced parton radiation. Considering that the mass of the bottom quark is much larger than the typical temperature of the medium, the momentum change during each interaction of the bottom quarks in the medium is relatively small, and it can be regarded as Brownian motion. Consequently, the momentum evolution of bottom quarks can be described using the Langevin equation:

$$\begin{array}{c}\hfill \frac{d\mathbf{p}}{dt}=-{\eta }_D\mathbf{p}+\mathit{\xi }+{\mathbf{f}}_g.\end{array}$$

(1)

On the right-hand side, the first two terms represent the drag and noise terms, which come from the elastic collisions with thermal light partons. The drag coefficient is defined as ${\eta }_D\left(p\right)=\kappa /\left(2T{E}_b\right)$, where the energy of the bottom quark is given by $E_b=\sqrt{m_b^2+p^2}$. The mass of the bottom quark is taken to be $m_b=4.75$ GeV. $\kappa$ represents the momentum–diffusion coefficient. It is related to the spatial diffusion coefficient $D_s$ via $D_s\kappa =2T^2$. In the noise term $\mathit{\xi }$, the time correlation and momentum dependence are both neglected for simplicity, where $\mathit{\xi }$ is treated as a Gaussian-shaped white noise that satisfies the following relation:

$$\begin{array}{c}\hfill \langle {\xi }^i\left(t\right){\xi }^j\left(t^{\prime }\right)\rangle =\kappa {\delta }^{ij}\delta (t-t^{\prime }).\end{array}$$

(2)

The third term ${\mathbf{f}}_g=-d{\mathbf{p}}_g/dt$, with ${\mathbf{p}}_g$ being the momentum of the emitted gluon, represents the recoil force on the bottom quark from the emitted gluon. The number of emitted gluons in a small time interval $t\sim t+\mathsf{\Delta }t$ is [46] are represented by

$$\begin{array}{c}\hfill P_{\mathrm{rad}}(t,\mathsf{\Delta }t)=\langle N_g(t,\mathsf{\Delta }t)\rangle =\mathsf{\Delta }t\int dxd{k}_T^2\frac{d{N}_g}{dxd{k}_T^{2}dt}.\end{array}$$

(3)

Here, $x=E_g/E_b$ represents the ratio of the energy carried by the gluons that are radiated from the bottom quarks. $d{N}_g/dxd{k}_T^{2}dt$ is the the spectrum of an emitted gluon from a higher twist calculation [22,23]. $k_T$ is the transverse momentum of the gluon. The position of a heavy quark is updated in each time step as $\mathbf{x}(t+\mathsf{\Delta }t)=\mathbf{x}\left(t\right)+\mathbf{p}/E_b·\mathsf{\Delta }t$.

Heavy quarks are randomly generated based on the initial spatial and momentum distributions, and they propagate in the QGP with an energy loss that is described by Equation (1). When bottom quarks diffuse and move into certain regions, wherein the QGP local temperature is low, heavy quarks undergo hadronization by combining with light quarks to form B mesons, or by combining with an anti-heavy quark to form quarkonium. In the high-momentum region, the production of B mesons from bottom quarks is predominantly through fragmentation processes, while in the intermediate- and low-momentum regions, it is mainly through the coalescence process. In this study, we primarily focused on bottom quarks with a transverse momentum of $p_T\le 15$ GeV/c. Therefore, we utilized a coalescence model to describe the hadronization process of bottom quarks into B mesons as follows:

$$\begin{array}{c}\hfill \frac{d{N}_M}{d{\mathbf{p}}_M}=\int \frac{d{\mathbf{p}}_1}{{\left(2\pi \right)}^3}\frac{d{\mathbf{p}}_2}{{\left(2\pi \right)}^3}\frac{d{N}_1}{d{\mathbf{p}}_1}\frac{d{N}_2}{d{\mathbf{p}}_2}{f}_M^W\left({\mathbf{q}}_r\right){\delta }^{\left(3\right)}({\mathbf{p}}_M-{\mathbf{p}}_1-{\mathbf{p}}_2),\end{array}$$

(4)

The momentum distribution of the B meson $d{N}_M/d{\mathbf{p}}_M$ is proportional to the distributions of the bottom quarks $d{N}_1/d{\mathbf{p}}_1$ and also to the thermal light quarks $d{N}_2/d{\mathbf{p}}_2$. The heavy quark distribution is given by the Langevin equation, while thermal light quarks are taken as a Fermi distribution. Their coalescence probability is determined by the Wigner function $f_M^W\left({\mathbf{q}}_r\right)$, which can be obtained via the Weyl transform of the B meson wave function. In principle, the complete Wigner function $f^W({\mathbf{q}}_r,{\mathbf{x}}_r)$ provides constraints on the relative distance and relative momentum that occurs between two particles in the formation of a bound state. The spatial constraint becomes crucial and can significantly reduce the coalescence probability when the two particles are rare in the QGP, as observed in the case of the charmonium coalescence process. However, in the case of a B meson composed of one heavy and one light anti-quark, with a plentiful number of light quarks in the QGP, we assume that the heavy quark can readily find a light quark in proximity, thus satisfying the spatial constraint. By integrating over the spatial part of the complete Wigner function, we simplify it to a normalized Gaussian function $A_{0}exp(-q_r^2{\sigma }^2)$, where $A_0$ represents the normalization factor. The width of the Gaussian function is connected with the root–mean–square radius of the B meson ${\sigma }^2=\frac{4}{3}\frac{{(m_1+m_2)}^2}{m_1^2+m_2^2}{\langle r^2\rangle }_B$ [47], where the value is to be $\sqrt{{\langle r^2\rangle }_B}=0.43$ fm [45]. $m_1$ is the bottom quark mass, while the thermal mass of the light quark is approximated to be $m_2=0.3$ GeV, which is used in the coalescence process. ${\mathbf{q}}_r=(E_2^{\mathrm{cm}}{\mathbf{p}}_1^{\mathrm{cm}}-E_1^{\mathrm{cm}}{\mathbf{p}}_2^{\mathrm{cm}})/(E_1^{\mathrm{cm}}+E_2^{\mathrm{cm}})$ is the relative momentum between the bottom quark and the light quark in the center of mass (COM) frame. ${\mathbf{p}}_1^{\mathrm{cm}}$ and ${\mathbf{p}}_2^{\mathrm{cm}}$ are the momenta of the bottom quark and the light quark in the COM frame. The delta function ensures the momentum conservation in the coalescence process ${\mathbf{p}}_M={\mathbf{p}}_1+{\mathbf{p}}_2$. As the phase transition between QGP and hadronic gas crossover at LHC energies, we performed the coalescence process at the critical temperature $T_c=150$ MeV. The time and spatial evolutions of bulk media have been well described with hydrodynamic equations. We employed the MUSIC package to provide the information for the hot medium at 5.02 TeV Pb–Pb collisions [27,48,49]. The local temperatures of the medium varied with coordinates and time, and these variations were then incorporated into the Langevin equation. After hadronization, the B mesons continue their diffusion in the hadronic gas with a different value of the diffusion coefficient, and they then decay into a non-prompt $J/\psi$ after a kinetic freeze out at a temperature of $T_{\mathrm{fo}}=120$ MeV.

Recently, through lattice QCD calculations, new calculations of the spatial diffusion coefficient of heavy quarks at different temperatures have been presented [20], and they were found to be smaller than the previous quenched lattice QCD [50,51] and recent phenomenological estimates [40,52,53,54]. This conclusion was also observed in other theoretical results [55,56,57]. This prompted us to re-examine the relationship between the experimental measurements of heavy quarks and the diffusion coefficient. In high temperature and momentum regimes, the diffusion coefficient can be calculated through perturbative QCD [58,59,60]. However, this calculation is not sufficient for simultaneously explaining the $R_{AA}$ and $v_2$ of open heavy flavor particles that are measured in nuclear collisions [61], thus suggesting that non-perturbative processes play an indispensable role in the temperature and momentum dependence of the diffusion coefficient. In Bayesian statistical analysis, a parameterized form of the diffusion coefficient is proposed, including a linear temperature dependence term and a perturbative QCD term such as $D_{s}2\pi T\propto A\left(p\right)(\alpha +\beta T)+(1-A\left(p\right))8\pi /(\widehat{q}/T^3)$ [40]. The first term represents the contribution from non-perturbative processes, while the second term represents the contribution from perturbative processes. $\widehat{q}$ is the heavy quark transport coefficient, which is calculated by the elastic scatterings between heavy and light quarks [62]. The spatial diffusion coefficient from lattice QCD calculations was found in the $p=0$ case. In this work, we introduce the following concise formula to consider the temperature and momentum dependence of the spatial diffusion coefficient:

$$\begin{array}{c}\hfill D_{s}2\pi T=[\alpha +\beta (\frac{T}{T_c}-1)]\times {\left(\frac{m_Q}{E_Q}\right)}^{\gamma }.\end{array}$$

(5)

The temperature and momentum dependence were encoded in the terms $T/T_c$ and $m_Q/E_Q$, where $m_Q$ and $E_Q=\sqrt{m_Q^2+p_Q^2}$ are the mass and energy of heavy quarks, respectively. The parameter $\alpha$ represents the value of $D_{s}2\pi T$ at a critical temperature, where the momentum is $p=0$. The parameters $\beta$ and $\gamma$ control the degree of temperature and momentum dependence. In hadronic gas, the coupling strength between B and D mesons with the medium becomes much smaller. Their contribution on the $R_{AA}$ and $v_2$ of open heavy flavor particles are limited. In hadronic gas, the mean value of the spatial diffusion coefficient of B mesons is approximated to be $D_s^{M}2\pi T=9$ before the kinetic freeze out of B mesons in $0.8T_c<T<T_c$ [63].

2.2. Deep Neural Networks

In the dynamical evolution of bottom quarks, our theoretical model produces a wide range of $R_{AA}$ and $v_2$ values for the non-prompt $J/\psi$ in 5.02 TeV Pb–Pb collisions by varying the shadowing factor (which is denoted as S) and the three parameters of $(\alpha ,\beta ,\gamma )$ in Equation (5). This dataset will serve as the training data for the CNN. Experimental measurements were conducted to determine the nuclear modification factor $R_{AA}\left(p_T\right)$ for the non-prompt $J/\psi$ in three different centralities, as well as the elliptic flow coefficient of $v_2\left(p_T\right)$. In our approach, we treat the three centralities of $R_{AA}\left(p_T\right)$ as separate channels, while $v_2\left(p_T\right)$ acts as an additional channel in the input layer of the CNN. The output layer of the CNN incorporates the corresponding parameter values in the diffusion coefficient and the shadowing factor, which are considered labels for the input data. Theoretical calculations based on the LICM model establish a mapping relationship between the parameter combination values $(S,\alpha ,\beta ,\gamma )$ and the experimental observables $(R_{AA},v_2)$. Figure 1 provides a graphical representation of the network structure, which consists of four hidden layers, including the average pooling layers and fully connected layers. The ReLU activation function was chosen for these layers. The three output channels related to the diffusion coefficient parameters were projected to have positive values, while the channel associated with the shadowing effect was projected within the range of 0–1 when using the Sigmoid function.

To generate the training dataset, we randomly selected parameter values within the regions specified in

Table 1. The samples of the parameters used in the training dataset.

Parameters	Region
Shadow factor	$S\in [0.6,1.0]$
$D_{s}2\pi T$	$\alpha \in [2.0,6.0]$
with
T-dependence	$\beta \in [0.0,8.0]$
$p_T$-dependence	$\gamma \in [0.0,1.0]$

using the Langevin model. Due to the significant uncertainty surrounding the shadowing factor S in the model calculations, which can notably impact the final observables of the B meson, we considered the shadowing factor as a parameter to be optimized within the deep neural network. The values of the shadowing factor were constrained within the range of 0.6 to 1.0. For the 5.02 TeV Pb–Pb collisions, the values of S used to generate the training dataset were randomly selected from the range of 0.6 to 1.0. The spatial diffusion coefficient $D_{s}2\pi T$ at the critical temperature of $T_c$ was selected within the range of $2.0\le \alpha \le 6.0$, while the values of $\beta$ and $\gamma$ that characterized the temperature and momentum dependence were chosen within the ranges of $0\le \beta \le 8.0$ and $0\le \gamma \le 1.0$, respectively. We generated 20 K events as the training and validation datasets, where each event corresponded to one of the combinations of the parameters. The performance of the neural network was influenced by both the size of the training dataset and the structure of the deep neural network. This type of uncertainty has been examined in previous studies [64], and it will be partially investigated in this work by varying the size of the training datasets. However, such uncertainties were not as significant as those arising from the error bars of the experimental data points of heavy flavor particles that were used as the input for the CNN.

3. Results and Analyses

In the previous sections, we introduced the theoretical model to generate the training dataset for the CNN by varying the values of the parameters in the model. We plotted some events ($(R_{AA}\left(p_T\right),v_2\left(p_T\right))$), and these were randomly selected from one channel of the training datasets, as shown in Figure 2. The lines could cover the experimental data, which indicated that the training range of the model encompassed the distribution of the experimental data. The model was able to be effectively applied to analyze the experimental data. As the experimental data points about $R_{AA}$ were located in $p_T\ge 2$ GeV/c, and as $v_2$ was located in $p_T\ge 4$ GeV/c, we truncated the training data by only retaining the data with $p_T$ values above 4 GeV/c. This ensured that the shape of the training data aligned as closely as possible with the experimental data, thus making it easier to incorporate them into the input (please see Figure 2).

We partitioned 70% of the total datasets as the training data, while the remaining 30% of the datasets served as the validation data. By treating $(R_{AA},v_2)$ as the inputs for the CNN and the parameter values as labels, we could calculate the loss of the CNN for both the training and validation datasets. The learning curve shown in Figure 3 demonstrates that the loss decreased to below 5% after 250 training epochs. Notably, the loss of the CNN for the training datasets closely aligned with the loss observed when using the validation datasets. This indicated that the CNN model did not exhibit significant overfitting or underfitting.

In using the well-trained CNN model, we fed the experimental data points into the neural network. Considering the presence of error bars in the experimental data points (as depicted in Figure 4), we sampled within the error bars associated with the experimental data points. These samples were then utilized as inputs to the neural network, thus allowing us to account for the impact of experimental uncertainties on the diffusion coefficient parameters. To capture the influence of the experimental uncertainty, we randomly generated 10 K samples within the range of experimental error bars around the data points, as illustrated in Figure 4.

Consequently, we obtained 10 K different outputs, which are plotted in Figure 5. Each data point represents an individual event, with the corresponding values of $\alpha$ and $\beta$ represented on the x- and y-axes, respectively. As previously mentioned, $\alpha$ signifies the value of $D_{s}2\pi T$ at critical temperature and zero momentum, while $\beta$ represents its temperature dependence. The majority of events were concentrated within the region of $4\le \alpha \le 6.5$ and $0\le \beta \le 5.0$. Furthermore, we also plotted the values of $\gamma$, which helped with characterizing the transverse momentum dependence, as well as the shadowing factor S that arises from the cold nuclear matter effect shown in Figure 6. From the distributions, it was observed that the majority of the events fell within the range of $0.0\le \gamma \le 0.2$ and $0.75\le S\le 0.9$, as depicted in Figure 6. The distribution of the events in the figure reflected the uncertainty in the parameter values that resulted from the experimental data errors, as well as the neural network structure.

Based on the distributions obtained from the CNN outputs, we can directly extract the mean values of these parameters, which are presented in

Table 2. The mean values and the standard deviation of the parameters were extracted using the experimental data points for the non-prompt $J/\psi$$R_{AA}\left(p_T\right)$ and $v_2\left(p_T\right)$ in the 5.02 TeV Pb–Pb Collisions.

Parameters	Mean Values	Standard Deviation
Shadow factor	$\langle S\rangle =0.82$	0.050
$D_{s}2\pi T$	$\langle \alpha \rangle =4.87$	0.90
with
T-dependence	$\langle \beta \rangle =4.16$	2.32
$p_T$-dependence	$\langle \gamma \rangle =0.058$	0.12

. It is worth noting that the value of $\alpha$ remains larger than the results obtained from the lattice QCD calculations at $T=T_c$ [20], and these are consistent with the values given by previous model calculations. Additionally, the temperature dependence ($\beta$) was found to be strong, while the momentum dependence ($\gamma$) was relatively weak.

4. Summary

In this study, a convolutional neural network was employed to extract the temperature and momentum dependence in the spatial diffusion coefficient of heavy quarks with the experimental data obtained from the non-prompt $J/\psi$ decays in B mesons. The Langevin equation was utilized to describe the dynamical evolution of the heavy quarks in the QGP and the B mesons in the hadronic gas. Additionally, the instantaneous coalescence model was used to describe the hadronization process from the bottom quarks to the B mesons. By taking different values for the shadowing factor and diffusion coefficient, the nuclear modification factors and elliptic flows of the non-prompt $J/\psi$ in multiple centralities of 5.02 TeV Pb–Pb collisions were generated. The CNN model was trained under supervision with model calculations. To extract the values of the diffusion coefficient, we sampled $(R_{AA},v_2)$ from the experimental data points along with their error bars, which were then used as inputs for the CNN. By doing so, the corresponding values of the diffusion coefficient and shadowing factor were obtained concurrently. The dispersion in the diffusion coefficient values can be partially attributed to the uncertainties present in the experimental data. The mean values of the diffusion coefficient were also extracted. This research contributes to the understanding of the heavy quark diffusion coefficient through a data-driven analysis approach.