#!/usr/bin/env python
# coding: utf-8

# # rf702_efficiencyfit_2D
# Special pdf's: unbinned maximum likelihood fit of an efficiency eff(x) function
# to a dataset D(x,cut), cut is a category encoding a selection whose efficiency as function
# of x should be described by eff(x)
# 
# 
# 
# 
# **Author:**  Clemens Lange, Wouter Verkerke (C++ version)  
# <i><small>This notebook tutorial was automatically generated with <a href= "https://github.com/root-project/root/blob/master/documentation/doxygen/converttonotebook.py">ROOTBOOK-izer</a> from the macro found in the ROOT repository  on Wednesday, April 17, 2024 at 11:19 AM.</small></i>

# In[1]:


import ROOT


flat = False


# Construct efficiency function e(x,y)
# -----------------------------------------------------------------------

# Declare variables x,mean, with associated name, title, value and allowed
# range

# In[2]:


x = ROOT.RooRealVar("x", "x", -10, 10)
y = ROOT.RooRealVar("y", "y", -10, 10)


# Efficiency function eff(x;a,b)

# In[3]:


ax = ROOT.RooRealVar("ax", "ay", 0.6, 0, 1)
bx = ROOT.RooRealVar("bx", "by", 5)
cx = ROOT.RooRealVar("cx", "cy", -1, -10, 10)

ay = ROOT.RooRealVar("ay", "ay", 0.2, 0, 1)
by = ROOT.RooRealVar("by", "by", 5)
cy = ROOT.RooRealVar("cy", "cy", -1, -10, 10)

effFunc = ROOT.RooFormulaVar(
    "effFunc", "((1-ax)+ax*cos((x-cx)/bx))*((1-ay)+ay*cos((y-cy)/by))", [ax, bx, cx, x, ay, by, cy, y]
)


# Acceptance state cut (1 or 0)

# In[4]:


cut = ROOT.RooCategory("cut", "cutr", {"accept": 1, "reject": 0})


# Construct conditional efficiency pdf E(cut|x,y)
# ---------------------------------------------------------------------------------------------

# Construct efficiency pdf eff(cut|x)

# In[5]:


effPdf = ROOT.RooEfficiency("effPdf", "effPdf", effFunc, cut, "accept")


# Generate data(x,y,cut) from a toy model
# -------------------------------------------------------------------------------

# Construct global shape pdf shape(x) and product model(x,cut) = eff(cut|x)*shape(x)
# (These are _only_ needed to generate some toy MC here to be used later)

# In[6]:


shapePdfX = ROOT.RooPolynomial("shapePdfX", "shapePdfX", x, [0 if flat else -0.095])
shapePdfY = ROOT.RooPolynomial("shapePdfY", "shapePdfY", y, [0 if flat else +0.095])
shapePdf = ROOT.RooProdPdf("shapePdf", "shapePdf", [shapePdfX, shapePdfY])
model = ROOT.RooProdPdf("model", "model", {shapePdf}, Conditional=({effPdf}, {cut}))


# Generate some toy data from model

# In[7]:


data = model.generate({x, y, cut}, 10000)


# Fit conditional efficiency pdf to data
# --------------------------------------------------------------------------

# Fit conditional efficiency pdf to data

# In[8]:


effPdf.fitTo(data, ConditionalObservables={x, y}, PrintLevel=-1)


# Plot fitted, data efficiency
# --------------------------------------------------------

# Make 2D histograms of all data, data and efficiency function

# In[9]:


hh_data_all = data.createHistogram("hh_data_all", x, Binning=8, YVar=dict(var=y, Binning=8))
hh_data_sel = data.createHistogram("hh_data_sel", x, Binning=8, YVar=dict(var=y, Binning=8), Cut="cut==cut::accept")
hh_eff = effFunc.createHistogram("hh_eff", x, Binning=50, YVar=dict(var=y, Binning=50))


# Some adjustsment for good visualization

# In[10]:


hh_data_all.SetMinimum(0)
hh_data_sel.SetMinimum(0)
hh_eff.SetMinimum(0)
hh_eff.SetLineColor(ROOT.kBlue)


# Draw all frames on a canvas

# In[11]:


ca = ROOT.TCanvas("rf702_efficiency_2D", "rf702_efficiency_2D", 1200, 400)
ca.Divide(3)
ca.cd(1)
ROOT.gPad.SetLeftMargin(0.15)
hh_data_all.GetZaxis().SetTitleOffset(1.8)
hh_data_all.Draw("lego")
ca.cd(2)
ROOT.gPad.SetLeftMargin(0.15)
hh_data_sel.GetZaxis().SetTitleOffset(1.8)
hh_data_sel.Draw("lego")
ca.cd(3)
ROOT.gPad.SetLeftMargin(0.15)
hh_eff.GetZaxis().SetTitleOffset(1.8)
hh_eff.Draw("surf")

ca.SaveAs("rf702_efficiency_2D.png")


# Draw all canvases 

# In[12]:


from ROOT import gROOT 
gROOT.GetListOfCanvases().Draw()