Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[RF] Update SBI tutorials with improved plot styling and labels #16692

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 64 additions & 25 deletions tutorials/roofit/rf615_simulation_based_inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@
from sklearn.neural_network import MLPClassifier

# The samples used for training the classifier in this tutorial / rescale for more accuracy
n_samples = 1000
n_samples = 10000

# Kills warning messages
ROOT.RooMsgService.instance().setGlobalKillBelow(ROOT.RooFit.WARNING)
Expand Down Expand Up @@ -85,6 +85,15 @@ def morphing(setting):
morph = ROOT.RooWrapperPdf("morph", "morph", morph_func, True)
workspace.Import(morph, Silence=True)

# Uncomment to see input plots for the first dimension (you might need to increase the morphed samples)
# f1 = x_var.frame(Title="linear morphing;x;pdf", Range=(-4, 8))
# for i in range(n_grid):
# workspace[f"histpdf{i}"].plotOn(f1)
# workspace["morph"].plotOn(f1, LineColor="r")
# c0 = ROOT.TCanvas()
# f1.Draw()
# input() # Wait for user input to proceed


# Class used in this case to demonstrate the use of SBI in Root
class SBI:
Expand Down Expand Up @@ -139,15 +148,16 @@ def train_classifier(self):
# Define the "observed" data in a workspace
def build_ws(mu_observed, sigma):
# using a workspace for easier processing inside the class
workspace = ROOT.RooWorkspace()
workspace.factory(f"Gaussian::gauss(x[-5,15], mu[0,4], {sigma})")
workspace.factory("Uniform::uniform(x)")
workspace.Print("v")
obs_data = workspace["gauss"].generate(workspace["x"], n_samples)
ws = ROOT.RooWorkspace()
ws.factory(f"Gaussian::gauss(x[-5,15], mu[0,4], {sigma})")
ws.factory("Uniform::uniform(x)")
ws["mu"].setVal(mu_observed)
ws.Print("v")
obs_data = ws["gauss"].generate(ws["x"], 1000)
obs_data.SetName("obs_data")
workspace.Import(obs_data, Silence=True)
ws.Import(obs_data, Silence=True)

return workspace
return ws


# The "observed" data
Expand Down Expand Up @@ -201,30 +211,59 @@ def learned_likelihood_ratio(x, mu):
ROOT.SetOwnership(nll_morph, True)

# Plot the negative logarithmic summed likelihood
frame1 = mu_var.frame(Title="NLL of SBI vs. Morphing", Range=(1.5, 2.5))
nll_gauss.plotOn(frame1, LineColor="g", ShiftToZero=True, Name="gauss")
nllr_learned.plotOn(frame1, LineColor="r", LineStyle="--", ShiftToZero=True, Name="learned")
frame1 = mu_var.frame(Title="NLL of SBI vs. Morphing;mu;NLL", Range=(2.2, 2.8))
nllr_learned.plotOn(frame1, LineColor="kP6Blue", ShiftToZero=True, Name="learned")
nll_gauss.plotOn(frame1, LineColor="kP6Blue+1", ShiftToZero=True, Name="gauss")
ROOT.RooAbsReal.setEvalErrorLoggingMode("Ignore") # Silence some warnings
nll_morph.plotOn(frame1, LineColor="c", ShiftToZero=True, Name="morphed")
nll_morph.plotOn(frame1, LineColor="kP6Blue+2", ShiftToZero=True, Name="morphed")
ROOT.RooAbsReal.setEvalErrorLoggingMode("PrintErrors")

# Plot the likelihood functions
frame2 = x_var.frame(Title="Learned vs analytical likelihhood function")
llhr_learned.plotOn(frame2, LineColor="r", LineStyle="--", Name="learned_ratio")
llhr_calc.plotOn(frame2, LineColor="g", Name="exact")

# Write the plots into one canvas
c = ROOT.TCanvas("rf615_simulation_based_inference", "rf615_simulation_based_inference", 800, 400)
c.Divide(2)
c.cd(1)
ROOT.gPad.SetLeftMargin(0.15)
frame1.GetYaxis().SetTitleOffset(1.8)
frame2 = x_var.frame(Title="Likelihood ratio r(x|#mu=2.5);x;p_{gauss}/p_{uniform}")
llhr_learned.plotOn(frame2, LineColor="kP6Blue", Name="learned_ratio")
llhr_calc.plotOn(frame2, LineColor="kP6Blue+1", Name="exact")

# Write the plots into one canvas to show, or into separate canvases for saving.
single_canvas = True

c = ROOT.TCanvas("", "", 1200 if single_canvas else 600, 600)
if single_canvas:
c.Divide(2)
c.cd(1)
ROOT.gPad.SetLeftMargin(0.15)
frame1.GetYaxis().SetTitleOffset(1.8)
frame1.Draw()
c.cd(2)
ROOT.gPad.SetLeftMargin(0.15)
frame2.GetYaxis().SetTitleOffset(1.8)

legend1 = ROOT.TLegend(0.43, 0.63, 0.8, 0.87)
legend1.SetFillColor(ROOT.kWhite)
legend1.SetLineColor(ROOT.kWhite)
legend1.SetTextSize(0.04)
legend1.AddEntry("learned", "learned (SBI)", "L")
legend1.AddEntry("gauss", "true NLL", "L")
legend1.AddEntry("morphed", "moment morphing", "L")
legend1.Draw()

if single_canvas:
c.cd(2)
ROOT.gPad.SetLeftMargin(0.15)
frame2.GetYaxis().SetTitleOffset(1.8)
else:
c.SaveAs("rf615_plot_1.png")
c = ROOT.TCanvas("", "", 600, 600)

frame2.Draw()

legend2 = ROOT.TLegend(0.53, 0.73, 0.87, 0.87)
legend2.SetFillColor(ROOT.kWhite)
legend2.SetLineColor(ROOT.kWhite)
legend2.SetTextSize(0.04)
legend2.AddEntry("learned_ratio", "learned (SBI)", "L")
legend2.AddEntry("exact", "true ratio", "L")
legend2.Draw()

if not single_canvas:
c.SaveAs("rf615_plot_2.png")

# Compute the minimum via minuit and display the results
for nll in [nll_gauss, nllr_learned, nll_morph]:
minimizer = ROOT.RooMinimizer(nll)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@
# Kills warning messages
ROOT.RooMsgService.instance().setGlobalKillBelow(ROOT.RooFit.WARNING)

n_samples_morph = 1000 # Number of samples for morphing
n_samples_morph = 10000 # Number of samples for morphing
n_bins = 4 # Number of 'sampled' Gaussians
n_samples_train = n_samples_morph * n_bins # To have a fair comparison

Expand Down Expand Up @@ -172,7 +172,7 @@ def train_classifier(self):
# Define the "observed" data in a workspace
def build_ws(mu_observed):
n_vars = len(mu_observed)
x_vars = [ROOT.RooRealVar(f"x{i}", f"x{i}", -2, 6) for i in range(n_vars)]
x_vars = [ROOT.RooRealVar(f"x{i}", f"x{i}", -5, 15) for i in range(n_vars)]
mu_vars = [ROOT.RooRealVar(f"mu{i}", f"mu{i}", mu_observed[i], 0, 4) for i in range(n_vars)]
gaussians = [ROOT.RooGaussian(f"gauss{i}", f"gauss{i}", x_vars[i], mu_vars[i], sigmas[i]) for i in range(n_vars)]
uniforms = [ROOT.RooUniform(f"uniform{i}", f"uniform{i}", x_vars[i]) for i in range(n_vars)]
Expand Down Expand Up @@ -257,14 +257,14 @@ def learned_likelihood_ratio(*args):

# Plot the learned and analytical summed negativelogarithmic likelihood
frame1 = mu_vars[0].frame(
Title="Negative logarithmic Likelihood",
Title="NLL of SBI vs. Morphing;#mu_{1};NLL",
Range=(mu_observed[0] - 1, mu_observed[0] + 1),
)
nll_gauss.plotOn(frame1, ShiftToZero=True, LineColor="g", Name="gauss")
nll_gauss.plotOn(frame1, ShiftToZero=True, LineColor="kP6Blue+1", Name="gauss")
ROOT.RooAbsReal.setEvalErrorLoggingMode("Ignore") # Silence some warnings
nll_morph.plotOn(frame1, ShiftToZero=True, LineColor="c", Name="morph")
nll_morph.plotOn(frame1, ShiftToZero=True, LineColor="kP6Blue+2", Name="morph")
ROOT.RooAbsReal.setEvalErrorLoggingMode("PrintErrors")
nllr_learned.plotOn(frame1, LineColor="r", ShiftToZero=True, LineStyle="--", Name="learned")
nllr_learned.plotOn(frame1, LineColor="kP6Blue", ShiftToZero=True, Name="learned")


# Declare a helper function in ROOT to dereference unique_ptr
Expand All @@ -281,23 +281,51 @@ def learned_likelihood_ratio(*args):
lhr_calc_final.recursiveRedirectServers(norm_set)

# Plot the likelihood ratio functions
frame2 = x_vars[0].frame(Title="Learned vs analytical likelihood function")
lhr_learned.plotOn(frame2, LineColor="r", LineStyle="--", Name="learned")
lhr_calc_final.plotOn(frame2, LineColor="g", Name="analytical")

c = ROOT.TCanvas(
"rf617_simulation_based_inference_multidimensional", "rf617_simulation_based_inference_multidimensional", 800, 400
)
c.Divide(2)
c.cd(1)
ROOT.gPad.SetLeftMargin(0.15)
frame1.GetYaxis().SetTitleOffset(1.8)
frame2 = x_vars[0].frame(Title="Likelihood ratio r(x_{1}|#mu_{1}=2.5);x_{1};p_{gauss}/p_{uniform}")
lhr_learned.plotOn(frame2, LineColor="kP6Blue", Name="learned_ratio")
lhr_calc_final.plotOn(frame2, LineColor="kP6Blue+1", Name="exact")

# Write the plots into one canvas to show, or into separate canvases for saving.
single_canvas = True

c = ROOT.TCanvas("", "", 1200 if single_canvas else 600, 600)
if single_canvas:
c.Divide(2)
c.cd(1)
ROOT.gPad.SetLeftMargin(0.15)
frame1.GetYaxis().SetTitleOffset(1.8)
frame1.Draw()
c.cd(2)
ROOT.gPad.SetLeftMargin(0.15)
frame2.GetYaxis().SetTitleOffset(1.8)

legend1 = ROOT.TLegend(0.43, 0.63, 0.8, 0.87)
legend1.SetFillColor(ROOT.kWhite)
legend1.SetLineColor(ROOT.kWhite)
legend1.SetTextSize(0.04)
legend1.AddEntry("learned", "learned (SBI)", "L")
legend1.AddEntry("gauss", "true NLL", "L")
legend1.AddEntry("morphed", "moment morphing", "L")
legend1.Draw()

if single_canvas:
c.cd(2)
ROOT.gPad.SetLeftMargin(0.15)
frame2.GetYaxis().SetTitleOffset(1.8)
else:
c.SaveAs("rf617_plot_1.png")
c = ROOT.TCanvas("", "", 600, 600)

frame2.Draw()

legend2 = ROOT.TLegend(0.53, 0.73, 0.87, 0.87)
legend2.SetFillColor(ROOT.kWhite)
legend2.SetLineColor(ROOT.kWhite)
legend2.SetTextSize(0.04)
legend2.AddEntry("learned_ratio", "learned (SBI)", "L")
legend2.AddEntry("exact", "true ratio", "L")
legend2.Draw()

if not single_canvas:
c.SaveAs("rf617_plot_2.png")


# Use ROOT's minimizer to compute the minimum and display the results
for nll in [nll_gauss, nllr_learned, nll_morph]:
Expand Down
50 changes: 31 additions & 19 deletions tutorials/roofit/rf618_mixture_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,17 +112,13 @@


# Define functions to compute the learned likelihood.
def calculate_likelihood_xgb(m4l_arr):
def calculate_likelihood_xgb(m4l_arr: np.ndarray) -> np.ndarray:
prob = model_xgb.predict_proba(m4l_arr.T)[:, 0]
return (1 - prob) / prob


llh = ROOT.RooFit.bindFunction(f"llh", calculate_likelihood_xgb, m4l)

# Plot the likelihood
frame1 = m4l.frame(Title=f"Likelihood", Range=(0, 200))
llh.plotOn(frame1, ShiftToZero=False)

# Number of signals and background
n_signal = results["higgs"]["weight"].sum()
n_back = results["zz"]["weight"].sum()
Expand All @@ -138,7 +134,7 @@ def weight_signal(mu):


# Define the likelihood ratio accordingly to mixture models
def likelihood_ratio(llr, mu):
def likelihood_ratio(llr: np.ndarray, mu: np.ndarray) -> np.ndarray:

m = 2

Expand All @@ -159,7 +155,12 @@ def likelihood_ratio(llr, mu):
mu_var = ROOT.RooRealVar("mu", "mu", 0.1, 5)

nll_ratio = ROOT.RooFit.bindFunction(f"nll", likelihood_ratio, llh, mu_var)
pdf_learned = ROOT.RooWrapperPdf("learned_pdf", "learned_pdf", nll_ratio, True)
pdf_learned = ROOT.RooWrapperPdf("learned_pdf", "learned_pdf", nll_ratio, selfNormalized=True)

# Plot the likelihood
frame1 = m4l.frame(Title="Likelihood ratio r(m_{4l}|#mu=1);m_{4l};p(#mu=1)/p(#mu=0)", Range=(80, 170))
# llh.plotOn(frame1, ShiftToZero=False, LineColor="kP6Blue")
nll_ratio.plotOn(frame1, ShiftToZero=False, LineColor="kP6Blue")

n_pred = ROOT.RooFormulaVar("n_pred", f"{n_back} + mu * {n_signal}", [mu_var])
pdf_learned_extended = ROOT.RooExtendPdf("final_pdf", "final_pdf", pdf_learned, n_pred)
Expand All @@ -169,21 +170,32 @@ def likelihood_ratio(llr, mu):
nll = pdf_learned_extended.createNLL(data, Extended=True)

# Plot the nll computet by the mixture model
frame2 = mu_var.frame(Title="NLL")
nll.plotOn(frame2)

# Write the plots into one canvas
c = ROOT.TCanvas("rf618_sbi_higgs", "rf618_sbi_higgs", 800, 400)
c.Divide(2)
c.cd(1)
ROOT.gPad.SetLeftMargin(0.15)
frame1.GetYaxis().SetTitleOffset(1.8)
frame2 = mu_var.frame(Title="NLL sum;#mu (signal strength);#Delta NLL", Range=(0.5, 4))
nll.plotOn(frame2, ShiftToZero=True, LineColor="kP6Blue")

# Write the plots into one canvas to show, or into separate canvases for saving.
single_canvas = True

c = ROOT.TCanvas("", "", 1200 if single_canvas else 600, 600)
if single_canvas:
c.Divide(2)
c.cd(1)
ROOT.gPad.SetLeftMargin(0.15)
frame1.GetYaxis().SetTitleOffset(1.8)
frame1.Draw()
c.cd(2)
ROOT.gPad.SetLeftMargin(0.15)
frame2.GetYaxis().SetTitleOffset(1.8)

if single_canvas:
c.cd(2)
ROOT.gPad.SetLeftMargin(0.15)
frame2.GetYaxis().SetTitleOffset(1.8)
else:
c.SaveAs("rf618_plot_1.png")
c = ROOT.TCanvas("", "", 600, 600)

frame2.Draw()

if not single_canvas:
c.SaveAs("rf618_plot_2.png")

# Compute the minimum via minuit and display the results
minimizer = ROOT.RooMinimizer(nll)
Expand Down
Loading