Skip to content

Accelerating RooFit with GPUs

Fitting data distributions with analytical models is a common task in CMS analyses. Depending on the size of the dataset and the number of free parameters in the model, the fitting can take a significant amount of time and become a bottleneck in the analysis.

The standard tool to perform such fits is RooFit, distributed with the ROOT framework. Since ROOT version 6.26, RooFit supports GPU-accelerated fitting using the CUDA backend, which allows you to speed up the fitting process by up to an order of magnitude.

The Purdue Analysis Facility supports this feature, allowing users to leverage available GPU resources to speed up their RooFit code. The feature is supported in both Jupyter Notebooks and Terminals, and for both the C++ ROOT interface and PyROOT.

Prerequisites

  1. Start your AF session with a GPU.

  2. Load the LCG view with the CUDA-enabled ROOT build. LCG "releases" and "views" are software stacks distributed by CERN.

    • If using a Jupyter Notebook: simply select the LCG_106b_cuda kernel.

    • If using a Terminal, run the following command:

      source /cvmfs/sft.cern.ch/lcg/views/LCG_106b_cuda/x86_64-el8-gcc11-opt/setup.sh

Warning

The CUDA-enabled ROOT build is currently available only via the LCG software stack. It is not available in other kernels, including the global Pixi environment.

The only supported ROOT version at the moment is 6.32.08.

Enabling the CUDA backend in RooFit

To enable the CUDA backend in RooFit, the only thing you need to do is pass the rt.RooFit.EvalBackend.Cuda() argument to the fitTo() command (in C++: RooFit::EvalBackend("cuda")).

Below is an example of code that uses the CUDA backend for fitting a Z-boson mass spectrum:

import ROOT as rt

inputfile = "workspace_ggh_All_Zfit_no_e_cut_UL_calib_cat5.root"

rt.EnableImplicitMT()

file = rt.TFile.Open(inputfile)

canvas = rt.TCanvas()
canvas.cd()

mass =  rt.RooRealVar("mh_ggh","mass (GeV)",100,85,99)
frame = mass.frame()

# Breit Wigner
bwWidth = rt.RooRealVar("bwz_Width" , "widthZ", 2.5, 0, 30)
bwmZ = rt.RooRealVar("bwz_mZ" , "mZ", 91.2, 90, 92)
sigma = rt.RooRealVar("sigma" , "sigma", 2, 0.0, 5.0)
bwWidth.setConstant(True)
model1_1 = rt.RooBreitWigner("bwz", "BWZ",mass, bwmZ, bwWidth)

# Double Sided Crystal Ball
mean = rt.RooRealVar("mean" , "mean", 0, -10, 10) # mean is mean relative to BW
sigma = rt.RooRealVar("sigma" , "sigma", 2, .2, 4.0)
alpha1 = rt.RooRealVar("alpha1" , "alpha1", 2, 0.01, 45)
n1 = rt.RooRealVar("n1" , "n1", 10, 0.01, 185)
alpha2 = rt.RooRealVar("alpha2" , "alpha2", 2.0, 0.01, 65)
n2 = rt.RooRealVar("n2" , "n2", 25, 0.01, 385)
model1_2 = rt.RooCrystalBall("dcb","dcb",mass, mean, sigma, alpha1, n1, alpha2, n2)

mass.setBins(10000,"cache") # cache is repre of the variable only used in FFT
mass.setBins(200) # bin to 200 bins otherwise, fitting with FFT conv is gonna take forever
mass.setMin("cache",50.5)
mass.setMax("cache",130.5)
model1 = rt.RooFFTConvPdf("BWxDCB", "BWxDCB", mass, model1_1, model1_2)

# Exponential background
coeff = rt.RooRealVar("coeff", "coeff", 0.001, 0.000001, 1.0)
shift = rt.RooRealVar("shift", "Offset", 91, 75, 105)
shifted_mass = rt.RooFormulaVar("shifted_mass", "@0-@1", rt.RooArgList(mass, shift))
model2 = rt.RooExponential("bkg", "bkg", shifted_mass, coeff)

sigfrac = rt.RooRealVar("sigfrac", "sigfrac", 0.99, 0, 1.0)
model = rt.RooAddPdf("model3", "model3", [model1, model2],sigfrac)

data = file.w.data("data_Zfit_no_e_cut_UL_calib_cat5")

model.fitTo(data, rt.RooFit.Save(), rt.RooFit.EvalBackend.Cuda()) #GPU

data.plotOn(frame)
model.plotOn(frame, rt.RooFit.LineColor(rt.kRed))
model.plotOn(frame, rt.RooFit.Components("BWxDCB"),rt.RooFit.LineColor(rt.kBlue))
model.plotOn(frame, rt.RooFit.Components("bkg"),rt.RooFit.LineColor(rt.kGreen))


frame.Draw()
canvas.Update()
canvas.Draw()

To run this code, you can download the input file workspace_ggh_All_Zfit_no_e_cut_UL_calib_cat5.root from https://cernbox.cern.ch/s/zKjJHZxRbDkADPf.