rf618_mixture_models.py

Go to the documentation of this file.

## \file
## \ingroup tutorial_roofit_main
## \notebook
## Use of mixture models in RooFit.
##
## This tutorial shows, how to use mixture models for Likelihood Calculation in ROOT. Instead of directly
## calculating the likelihood we use simulation based inference (SBI) as shown in tutorial 'rf615_simulation_based_inference.py'.
## We train the classifier to discriminate between samples from an background hypothesis here the zz samples and a target
## hypothesis, here the higgs samples. The data preparation is based on the tutorial 'df106_HiggsToFourLeptons.py'.
##
## An introduction to mixture models can be found here https://arxiv.org/pdf/1506.02169.
##
## A short summary:
## We assume the whole probability distribution can be written as a mixture of several components, i.e.
## $$p(x|\theta)= \sum_{c}w_{c}(\theta)p_{c}(x|\theta)$$
## We can write the likelihood ratio in terms of pairwise classification problems
## \begin{align*}
##  \frac{p(x|\mu)}{p(x|0)}&= \frac{\sum_{c}w_{c}(\mu)p_{c}(x|\mu)}{\sum_{c'}w_{c'}(0)p_{c'}(x|0)}\\
##  &=\sum_{c}\Bigg[\sum_{c'}\frac{w_{c'}(0)}{w_{c}(\mu)}\frac{p_{c'}(x|0)}{p_{c}(x|\mu)}\Bigg]^{-1},
## \end{align*}
## where mu is the signal strength, and a value of 0 corresponds to the background hypothesis. Using this decomposition, one is able to use the pairwise likelihood ratios.
##
## Since the only free parameter in our case is mu, the distributions are independent of this parameter and the dependence on the signal strength can be encoded into the weights.
## Thus, the subratios simplify dramatically since they are independent of theta and these ratios can be pre-computed and the classifier does
## not need to be parametrized.
##
## If you wish to see an analysis done with template histograms see 'hf001_example.py'.
##
## \macro_image
## \macro_code
## \macro_output
##
## \date September 2024
## \author Robin Syring
 
import ROOT
import os
import numpy as np
import xgboost as xgb
 
# Get Dataframe from tutorial df106_HiggsToFourLeptons.py
# Adjust the path if running locally
df = ROOT.RDataFrame("tree", ROOT.gROOT.GetTutorialDir().Data() + "/analysis/dataframe/df106_HiggsToFourLeptons.root")
 
# Initialize a dictionary to store counts and weight sums for each category
results = {}
 
 
# Extract the relevant columns once and avoid repeated calls
data_dict = df.AsNumpy(columns=["m4l", "sample_category", "weight"])
 
 
weights_dict = {
    name: data_dict["weight"][data_dict["sample_category"] == [name]].sum() for name in ("data", "zz", "other", "higgs")
}
 
# Loop over each sample category
for sample_category in ["data", "higgs", "zz", "other"]:
 
    weight_sum = weights_dict[sample_category]
 
    mask = data_dict["sample_category"] == sample_category
    # Normalize each weight
    weights = data_dict["weight"][mask]
    # Extract the weight_modified
    weight_modified = weights / weight_sum
 
    count = np.sum(mask)
 
    # Store the count and weight sum in the dictionary
    results[sample_category] = {
        "weight_sum": weight_sum,
        "weight_modified": weight_modified,
        "count": count,
        "weight": weights,
    }
 
 
# Extract the mass for higgs and zz
higgs_data = data_dict["m4l"][data_dict["sample_category"] == ["higgs"]]
zz_data = data_dict["m4l"][data_dict["sample_category"] == ["zz"]]
 
 
# Prepare sample weights
sample_weight_higgs = np.array([results["higgs"]["weight_modified"]]).flatten()
sample_weight_zz = np.array([results["zz"]["weight_modified"]]).flatten()
 
# Putting sample weights together in the same manner as the training data
sample_weight = np.concatenate([sample_weight_higgs, sample_weight_zz])
 
# For Training purposes we have to get rid of the negative weights, since xgb can't handle them
sample_weight[sample_weight < 0] = 1e-6
 
# Prepare the features and labels
X = np.concatenate((higgs_data, zz_data), axis=0).reshape(-1, 1)
y = np.concatenate([np.ones(len(higgs_data)), np.zeros(len(zz_data))])
 
# Train the Classifier to discriminate between higgs and zz
model_xgb = xgb.XGBClassifier(n_estimators=1000, max_depth=5, eta=0.2, min_child_weight=1e-6, nthread=1)
model_xgb.fit(X, y, sample_weight=sample_weight)
 
 
# Building a RooRealVar based on the observed data
m4l = ROOT.RooRealVar("m4l", "Four Lepton Invariant Mass", 0.0)
 
 
# Define functions to compute the learned likelihood.
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)
 
# Number of signals and background
n_signal = results["higgs"]["weight"].sum()
n_back = results["zz"]["weight"].sum()
 
 
# Define weight functions
def weight_back(mu):
    return n_back / (n_back + mu * n_signal)
 
 
def weight_signal(mu):
    return 1 - weight_back(mu)
 
 
# Define the likelihood ratio accordingly to mixture models
def likelihood_ratio(llr: np.ndarray, mu: np.ndarray) -> np.ndarray:
 
    m = 2
 
    w_0 = np.array([weight_back(0), weight_signal(0)])
    w_1 = np.array([weight_back(mu[0]), weight_signal(mu[0])])
 
    w = np.outer(w_1, 1.0 / w_0)
 
    p = np.ones((m, m, len(llr)))
    p[1, 0] = llr
    for i in range(m):
        for j in range(i):
            p[j, i] = 1.0 / p[i, j]
 
    return 1.0 / np.sum(1.0 / np.sum(np.expand_dims(w, axis=2) * p, axis=0), axis=0)
 
 
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, 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)
 
# Prepare the observed data set and NLL
data = ROOT.RooDataSet.from_numpy({"m4l": data_dict["m4l"][data_dict["sample_category"] == ["data"]]}, [m4l])
nll = pdf_learned_extended.createNLL(data, Extended=True)
 
# Plot the nll computet by the mixture model
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()
 
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)
minimizer.setErrorLevel(0.5)  # Adjust the error level in the minimization to work with likelihoods
minimizer.setPrintLevel(-1)
minimizer.minimize("Minuit2")
result = minimizer.save()
ROOT.SetOwnership(result, True)
result.Print()
 
del minimizer
del nll
del pdf_learned_extended
del n_pred
del llh
del nll_ratio