Introduction

Machine learning has become an indispensable tool across a broad range of scientific and technological disciplines, including image recognition, natural-language processing, recommendation systems, and speech analysis. Its main advantage is the ability to identify complex structures in high-dimensional datasets and to provide powerful tools for classification, regression, and clustering [1-6]. In nuclear physics, machine learning has been used to extract or predict key observables such as charge radii, nuclear masses, and binding energies [7-11]. In relativistic heavy-ion collisions, it is especially valuable for classification tasks, for example separating rare physics signals from large combinatorial backgrounds [12-17]. Representative algorithms include boosted decision trees (BDTs) [18, 19], deep neural networks [20-22], and ensemble methods [23], all of which can exploit multidimensional correlations more effectively than conventional straight-cut analyses. Independent of the specific algorithm, a reliable classifier requires representative signal and background samples. Whereas background samples can often be obtained directly from data, signal samples are usually produced with Monte Carlo simulations. Because detector simulations do not perfectly reproduce the data, training on mismatched signal samples can degrade classifier performance and bias the extracted physics observables, as illustrated in Appendix 7.

To address this challenge, we propose two correction methods: cumulative distribution function (CDF) mapping and shift-and-scale. CDF mapping systematically transforms a simulated feature distribution into the corresponding distribution measured in data without requiring an explicit parametric description of the shape, which makes it flexible and broadly applicable. The shift-and-scale procedure is more restrictive but simpler to implement, and it is well suited to cases in which the simulated and measured distributions already have similar shapes and differ mainly in location and width.

J/ψ is a classic probe of the medium created in relativistic heavy-ion collisions and is widely used to investigate the properties of the Quark–Gluon Plasma [24, 25]. Accurate yield measurements are therefore of continuing interest [26-34]. In this work, a BDT implemented with the eXtreme Gradient Boosting (XGBoost) framework [35] is used to separate the J/ψ signal from background in Ru+Ru and Zr+Zr collisions at GeV recorded by STAR at RHIC. This case study demonstrates both the implementation of the proposed correction procedures and their effectiveness in improving the robustness of machine-learning-based signal extraction. Section 2 describes the J/ψ reconstruction, Sections 3 and 4 present the construction of the training samples and the model training procedure, Sect. 5 discusses the self-consistency tests and results, and Sect. 6 summarizes the main conclusions.

J/ψ Reconstruction

In 200 GeV Ru+Ru and Zr+Zr collisions, J/ψ mesons are reconstructed through the dielectron decay channel [27, 34, 36], with electron¹ identification based on the Time-of-Flight (TOF) [37], the Time Projection Chamber (TPC) [38, 39], and the Barrel Electromagnetic Calorimeter (BEMC) [40] detectors. Loose preselection requirements on the particle-identification (PID) observables are first imposed to suppress the dominant background while retaining sufficient signal statistics for the subsequent machine-learning analysis.

The TOF detector measures the arrival time of a charged particle. Together with the collision start time and the momentum measured in the TPC, one obtains , where c is the speed of light and v is the particle velocity. Figure 1 shows the distribution of 1/β as a function of momentum p for charged particles. Because particles with different masses have different velocities at the same momentum, distinct PID bands appear in this representation. A preselection cut of is applied, shown as dashed lines in Fig. 1 and listed in Table 1.

Fig. 1Full-screen View

(Color online) Distribution of 1/β versus momentum for charged particles, where the dashed horizontal lines indicate the 1/β cuts for electron preselection

The TPC measures the specific ionization energy loss, dE/dx, which provides additional separation between electrons and hadrons. PID based on dE/dx is commonly expressed through the normalized deviation of the measured value from the expectation for a given particle species, denoted by nσ_e, nσπ, and so forth. For electrons, the nσ_e distribution is expected to be approximately Gaussian with mean 0 and width 1 after proper calibration. Figure 2 shows the nσ_e distribution as a function of p after imposing a tighter TOF requirement, , to improve the electron purity for visualization. The dashed curves indicate the momentum-dependent preselection criteria used before machine-learning training.

Fig. 2Full-screen View

(Color online) Distribution of nσ_e versus momentum for charged particles after applying the cut. The dashed curves correspond to the nσ_e cuts used for electron preselection

At high momentum, the BEMC provides additional PID through the ratio of the deposited electromagnetic energy to the track momentum. In practice, electron identification commonly uses the largest tower energy within the BEMC cluster, E₀, rather than the total cluster energy E_tot, because E₀ has better resolution and is less sensitive to the high detector occupancy in heavy-ion collisions. Figure 3 compares the E₀/p distributions for protons and electrons in the range 4<p_T<6 GeV/c. The electron sample, often referred to as photonic electrons, is obtained from photon conversions and neutral-meson Dalitz decays [41]. As expected, the electron distribution peaks at a larger E₀/p value than the proton distribution. The dashed vertical lines indicate the preselection requirement on E₀/p.

Fig. 3Full-screen View

(Color online) Distributions of E₀/p for electrons (circles) and protons (solid line) within 4<p_T<6 GeV/c. The vertical dashed lines correspond to the E₀/p cuts used for electron preselection

The use of the TOF, TPC, and BEMC preselection criteria depends on the track kinematics and on whether matching information is available in the TOF and BEMC, as summarized in Table 1. Electron and positron candidates that satisfy these requirements are paired to form J/ψ candidates, and the resulting invariant-mass distribution is shown in Fig. 4(a). The combinatorial background is estimated with the mixed-event technique [42]. After subtracting this component, the remaining invariant-mass distribution is fitted with two terms: a Crystal Ball function for the J/ψ signal and a linear term for the residual background arising mainly from semileptonic heavy-flavor decays and Drell–Yan production. The raw J/ψ yield is obtained by counting the electron–positron pairs in the mass window [2.91, 3.21] GeV/c² and subtracting the fitted residual-background contribution in the same interval. Approximately J/ψ candidates are reconstructed, with a signal significance of 103. The following sections show how machine learning further improves this significance.

Fig. 4Full-screen View

Invariant mass distributions of J/ψ candidates reconstructed via the dielectron channel, within and p_T > 0.2 GeV/c, for Ru+Ru and Zr+Zr collisions. Different panels correspond to different electron identification cuts. (a) With preselection cuts listed in Table 1; (b) With preselection and BDT cuts listed in Table 1; (c)With straight cuts listed in Table 2

Construction of Training Samples

For J/ψ reconstruction, the classification targets are electron–positron pairs that form J/ψ candidates. To determine whether a candidate is signal-like or background-like, a set of discriminating features must be defined. In principle, any observable that improves the separation without introducing physics bias can be used. In this work, the selected features are nσ_e, nσπ, E₀/p, and the distance of closest approach (DCA, defined as the minimum distance between a charged-particle trajectory and the collision vertex) for both daughter tracks, labeled as “tr1” and “tr2”.

The background training sample is selected from electron–positron pairs in the invariant-mass sidebands, namely (2.65, 2.85), (3.25, 3.5), and (3.8, 4.5) GeV/c², well separated from the J/ψ peak. The interval (3.5, 3.8) GeV/c² is excluded because it contains a small contribution from mesons. This sideband method assumes that the background in the sidebands is representative of the background beneath the J/ψ peak. Therefore, the features used for classification should not depend strongly on the pair invariant mass, which constitutes an important feature-selection criterion. Figure 5 shows the correlations between the selected features and the invariant mass for the background sample; no strong mass dependence is observed. As expected, nσ_e and nσπ are strongly correlated because both are derived from the same dE/dx measurement.

Fig. 5Full-screen View

(Color online) Correlations between all selected features and the reconstructed electron–positron pair invariant mass for the background training sample. Each cell represents the correlation coefficient between a pair of variables, with positive correlations shown in red and negative correlations in blue

The signal training sample is generated with Monte Carlo simulation. Electrons and positrons from J/ψ decays are propagated through a GEANT3-based [43] model of the STAR detector, and the simulated hits are embedded into real Ru+Ru and Zr+Zr events before reconstruction with the same software chain used for data. Inevitably, the simulation does not perfectly reproduce the data because detector conditions evolve with time and the heavy-ion environment has high occupancy. The quality of the simulation can be assessed by comparing the nσ_e, nσπ, E₀/p, and DCA distributions for electrons from simulated J/ψ decays with those for photonic electrons in data, as shown in Fig. 6. The DCA distributions agree reasonably well, whereas the nσ_e and nσπ distributions exhibit clear offsets and the E₀/p distribution shows a noticeable shape mismatch. If left uncorrected, these discrepancies can degrade the machine-learning performance and bias the extracted J/ψ yield.

Fig. 6Full-screen View

(Color online) Comparisons of feature distributions for electrons within GeV/c and from simulations, both with (open diamonds) and without (filled diamonds) corrections, to those from real data (filled circles). (a) nσ_e distributions; (b) nσπ distributions; (c) E₀/p distributions; (d) DCA distributions

Because the simulated and measured nσ_e and nσπ distributions have similar shapes, they can be corrected with a shift-and-scale transformation:(1)Here, μ_data (σ_data) and μ_simu (σ_simu) are the means (standard deviations) of the corresponding nσ_i distributions in data and simulation, respectively, with i=e or π. After this transformation, the simulated nσ_e and nσπ distributions agree well with those observed in data, as shown in Fig. 6(a) and (b).

For the E₀/p distribution, the mismatch between simulation and data is not limited to a simple offset or width difference, so shift-and-scale is insufficient. We therefore introduce a cumulative distribution function (CDF) mapping method. Let and , with , denote the probability density functions of a feature in simulation and data, respectively. The corresponding CDF isFor an event in simulation with feature value A, its cumulative probability isThe corrected value B is then obtained from the data distribution throughThis procedure defines an event-by-event mapping from the original value A to the corrected value B, such that the transformed simulation reproduces the data distribution. After CDF mapping, the corrected simulated E₀/p distribution (open diamonds) is consistent with the distribution measured in data (solid circles), as shown in Fig. 6(c). In practice, the PDFs are represented by discrete histograms, so the method is implemented numerically by constructing discrete CDFs and searching for the closest cumulative probability. When the simulation and data differ strongly in shape, this discrete matching can produce bin-to-bin fluctuations at the level of one histogram bin, leading to localized excesses in some bins and deficits in neighboring bins. This effect explains the statistical irregularities visible in Fig. 6(c). Because the classifier is trained on event-by-event continuous feature values rather than on the displayed histogram counts themselves, this binning artifact does not compromise the model training. It can be further reduced by decreasing the histogram bin width.

Compared with shift-and-scale, CDF mapping is more general but also more involved to implement. Both procedures act deterministically on each event and therefore retain the event-level correspondence among observables, which is essential for multivariate classifiers to benefit from joint feature information. For the same reason, the signal sample cannot be constructed by independently sampling one-dimensional feature distributions from data, because that procedure would destroy the physically relevant inter-feature dependence structure.

Model Training and Application

A BDT classifier implemented with XGBoost is trained to further improve the J/ψ signal significance. The signal and background samples are each divided into two equal subsets, one for training and the other for testing. After training, the model is applied to both subsets. For every electron–positron pair, the classifier returns an output between 0 and 1, with values close to 1 corresponding to signal-like candidates and values close to 0 corresponding to background-like candidates. The output distributions are shown in Fig. 7. The red band and red points denote the signal distributions for the training and testing samples, respectively, whereas the blue band and blue points show the corresponding background distributions. The strong separation between signal and background indicates good discriminating power, and the agreement between the training and testing samples shows that overtraining is not evident.

Fig. 7Full-screen View

(Color online) Model output distributions for both the training and testing samples. The red (blue) histogram represents the training signal (background) sample, while the red (blue) points are for the testing signal (background) sample

In the final analysis, a threshold on the BDT output is chosen so that pairs with scores above the threshold are retained as signal candidates and those below it are rejected as background. For a given threshold, the signal efficiency is defined as the fraction of signal events that survive the cut, and the background efficiency is defined analogously. Plotting the signal efficiency against the background efficiency while scanning the threshold from 0 to 1 yields the receiver operating characteristic (ROC) curve, where the horizontal and vertical axes correspond to the false-positive rate (FPR) and true-positive rate (TPR), respectively. The ROC curves for the training and testing samples are shown in Fig. 8. The area under the curve (AUC) quantifies the overall classification performance; values closer to 1 indicate better separation. The AUC obtained here is 0.91, demonstrating strong discriminating power. The nearly identical ROC curves and AUC values for the training and testing samples provide additional evidence that the model is not overfit.

Fig. 8Full-screen View

ROC curves for the training (dashed) and testing (solid) samples. The reported AUC values quantify the overall performance, while the diagonal line represents a random classifier for comparison

Model Validation via Self-Consistency Checks

Beyond establishing that the model is not overtrained, it is necessary to verify that the BDT efficiencies for both signal and background are correctly reproduced in real data. This validation is carried out through two self-consistency tests.

The first test examines the J/ψ yield corrected for the BDT efficiency as a function of the BDT threshold. Because the measured yield should not depend on the particular analysis threshold, the corrected yield is expected to remain constant within uncertainties. The J/ψ efficiency for a given BDT threshold is derived from the signal output distribution in Fig. 7, and the resulting efficiency curves are shown in Fig. 9 for several centrality intervals in Ru+Ru and Zr+Zr collisions. Here, centrality characterizes the geometric overlap of the two nuclei: 0–20% corresponds to the most central collisions, whereas 40–80% corresponds to more peripheral collisions. After applying the preselection criteria in Table 1 and varying the BDT threshold, the raw J/ψ counts are extracted from data with the procedure described in Sect. 2. Figure 10 then shows the efficiency-corrected J/ψ counts as a function of the BDT threshold for four centrality classes. Although the signal efficiency changes by roughly an order of magnitude as the threshold increases from 0 to 0.9, the corrected counts remain stable, with variations below 2%. This stability strongly suggests that the signal-efficiency estimate from the machine-learning model accurately represents the efficiency in data, thereby validating the feature-correction procedure described in Sect. 3.

Fig. 9Full-screen View

(Color online) BDT cut efficiencies for J/ψ with p_T>0.2 GeV/c in different centrality classes of 200 GeV Ru+Ru and Zr+Zr collisions

Fig. 10Full-screen View

(Color online) J/ψ counts, corrected by BDT cut efficiencies, as a function of BDT criterion in different centrality classes of 200 GeV Ru+Ru and Zr+Zr collisions

A more stringent validation compares the J/ψ significance as a function of the BDT threshold between the machine-learning expectation and the measurement in data, because this test depends on both the signal and background efficiencies. The expected J/ψ significance is calculated as(2)In Eq. 2, x is the BDT threshold, and and are the corresponding signal and background efficiencies. The quantities S(x) and B(x) denote the signal and background counts after the threshold is applied, while S₀ and B₀ are the counts at a reference threshold x₀. The choice of x₀ is arbitrary and is used only to normalize the prediction to the data at one point. The resulting expected significance is shown as the gray band in Fig. 11. It first increases and then decreases as the threshold is tightened from 0 to 0.9, reflecting the competition between improved background rejection and reduced signal efficiency. The prediction from the trained model is compared with the significance extracted directly from data, shown as open circles. For clarity, the two curves are aligned at a threshold of 0.5, namely x₀=0.5 in Eq. 2. Their consistency indicates that the training samples reproduce the behavior of both signal and background in data with sufficient accuracy for this analysis.

Fig. 11Full-screen View

(Color online) J/ψ significance as a function of BDT criterion. The black open circles represent the J/ψ significance directly obtained from real data while the black band shows the expected significance estimated from the machine learning model. The two distributions are aligned at BDT value of 0.5 as indicated by the black open square. The vertical dashed line at 0.72 marks the point where the expected significance reaches its maximum value of 176 as indicated by the horizontal dashed line. The red star represents the J/ψ significance of 173 obtained from real data at the same BDT criterion

The expected-significance curve from the trained model also provides a principled way to determine the optimal BDT threshold. Compared with scanning the data directly, an optimization based on the training samples is less sensitive to statistical fluctuations in the measured spectrum and therefore less prone to introducing threshold-selection bias, which is particularly important for low-statistics signals. As shown in Fig. 11, the J/ψ significance reaches its maximum at a BDT threshold of 0.72, which is therefore chosen as the working point. Figure 4(b) shows the raw signal extraction obtained with this optimal threshold. The corresponding significance is 173, representing a 68% improvement over the result with only the preselection criteria [Fig. 4(a)] and a 24% improvement over the result obtained with straight PID cuts [Fig. 4(c) and Tab. 2].

The same self-consistency tests are also useful for diagnosing problematic input features. If adding a given feature destroys the agreement between the model expectation and the data, that feature should be corrected or removed from the training. As an example, Fig. 12 in the Appendix shows the result obtained without CDF mapping for E₀/p. In that case, the efficiency-corrected signal counts deviate progressively from the expected flat behavior as the BDT threshold increases, reaching 29% at a threshold of 0.9. Likewise, the expected J/ψ significance becomes inconsistent with the significance measured in data, as shown in Fig. 13. These discrepancies provide direct evidence that the feature-correction procedures introduced in Sect. 3 are necessary for unbiased signal extraction.

Fig. 12Full-screen View

(Color online) J/ψ counts, corrected for the BDT cut efficiencies, as a function of BDT criterion, for scenarios that the signal sample is incorporated with (circles) or without (diamonds) the E₀/p correction

Fig. 13Full-screen View

(Color online) Expected J/ψ significance as a function of the BDT criterion from the machine learning model trained using the signal sample without the E₀/p correction. It is compared to the measured significance in real data (open circles). The two distributions are aligned at the BDT criterion of 0.5, indicated by the open square, to facilitate the comparison

Summary

This work develops two practical procedures, CDF mapping and shift-and-scale, to align simulated feature distributions with those measured in data while preserving the event-level multivariate information required for machine-learning classification. These methods address a common source of systematic bias in experimental analyses, namely the mismatch between simulated and real detector responses. Using J/ψ yield extraction in GeV Ru+Ru and Zr+Zr collisions as a case study, we show that the corrected training samples lead to stable efficiency-corrected yields and to a predicted significance consistent with the significance measured in data. The self-consistency tests therefore validate both the necessity and the effectiveness of the proposed corrections. The same strategy should be broadly useful in machine-learning applications throughout high-energy nuclear and particle physics whenever simulated signal samples do not fully reproduce the data.