Tutorial

This tutorial demonstrates the usage of the sweights package.

We will first cook up a toy model with a discriminat variable (invariant mass) and a control variable (decay time) and use it to generate some toy data.

Then we will use a fit to the invariant mass to obtain some component pdf estimates and use these to extract some weights which project out the signal only component in the decay time.

We will demonstrate both the classic sWeights and the Custom Ortogonal Weight functions (COWs) method. See arXiv:2112.04575 for more details.

Finally we will fit the weighted decay time distribution and correct the covariance matrix according to the description in arXiv:1911.01303.

[1]:
# external requirements
import os
import numpy as np
from scipy.stats import norm, expon, uniform
import matplotlib.pyplot as plt
from iminuit import Minuit
from iminuit.cost import ExtendedUnbinnedNLL
from iminuit.pdg_format import pdg_format
import boost_histogram as bh

# from this package
from sweights import SWeight # for classic sweights
from sweights import Cow     # for custom orthogonal weight functions
from sweights import cov_correct, approx_cov_correct # for covariance corrections
from sweights import kendall_tau # for an independence test
from sweights import plot_indep_scatter

# set a reproducible seed
np.random.seed(21011987)

Make the toy model and generate some data

[2]:
Ns = 5000
Nb = 5000
ypars = [Ns,Nb]

# mass
mrange = (0,1)
mu = 0.5
sg = 0.1
lb = 1
mpars = [mu,sg,lb]

# decay time
trange = (0,1)
tlb = 2
tpars = [tlb]

# generate the toy
def generate(Ns,Nb,mu,sg,lb,tlb,poisson=False,ret_true=False):

    Nsig = np.random.poisson(Ns) if poisson else Ns
    Nbkg = np.random.poisson(Nb) if poisson else Nb

    sigM = norm(mu,sg)
    bkgM = expon(mrange[0], lb)

    sigT = expon(trange[0], tlb)
    bkgT = uniform(trange[0],trange[1]-trange[0])

    # generate
    sigMflt = sigM.cdf(mrange)
    bkgMflt = bkgM.cdf(mrange)
    sigTflt = sigT.cdf(trange)
    bkgTflt = bkgT.cdf(trange)

    sigMvals = sigM.ppf( np.random.uniform(*sigMflt,size=Nsig) )
    sigTvals = sigT.ppf( np.random.uniform(*sigTflt,size=Nsig) )

    bkgMvals = bkgM.ppf( np.random.uniform(*bkgMflt,size=Nbkg) )
    bkgTvals = bkgT.ppf( np.random.uniform(*bkgTflt,size=Nbkg) )

    Mvals = np.concatenate( (sigMvals, bkgMvals) )
    Tvals = np.concatenate( (sigTvals, bkgTvals) )

    truth = np.concatenate( ( np.ones_like(sigMvals), np.zeros_like(bkgMvals) ) )

    if ret_true:
        return np.stack( (Mvals,Tvals,truth), axis=1 )
    else:
        return np.stack( (Mvals,Tvals), axis=1 )

toy = generate(Ns,Nb,mu,sg,lb,tlb,ret_true=True)

Plot the toy data and generating pdfs

[3]:
# useful function for plotting data points with error bars
def myerrorbar(data, ax, bins, range, wts=None, label=None, col=None):
    col = col or 'k'
    nh, xe = np.histogram(data,bins=bins,range=range)
    cx = 0.5*(xe[1:]+xe[:-1])
    err = nh**0.5
    if wts is not None:
        whist = bh.Histogram( bh.axis.Regular(bins,*range), storage=bh.storage.Weight() )
        whist.fill( data, weight = wts )
        cx = whist.axes[0].centers
        nh = whist.view().value
        err = whist.view().variance**0.5

    ax.errorbar(cx, nh, err, capsize=2,label=label,fmt=f'{col}o')

# define the mass pdf for plotting etc.
def mpdf(x, Ns, Nb, mu, sg, lb, comps=['sig','bkg']):

    sig  = norm(mu,sg)
    sigN = np.diff( sig.cdf(mrange) )

    bkg  = expon(mrange[0], lb)
    bkgN = np.diff( bkg.cdf(mrange) )

    tot = 0
    if 'sig' in comps: tot += Ns * sig.pdf(x) / sigN
    if 'bkg' in comps: tot += Nb * bkg.pdf(x) / bkgN

    return tot

# define time pdf for plotting etc.
def tpdf(x, Ns, Nb, tlb, comps=['sig','bkg']):

    sig  = expon(trange[0],tlb)
    sigN = np.diff( sig.cdf(trange) )

    bkg  = uniform(trange[0],trange[1]-trange[0])
    bkgN = np.diff( bkg.cdf(trange) )

    tot = 0
    if 'sig' in comps: tot += Ns * sig.pdf(x) / sigN
    if 'bkg' in comps: tot += Nb * bkg.pdf(x) / bkgN

    return tot

# define plot function
def plot(toy, draw_pdf=True):

    nbins = 50

    fig, ax = plt.subplots(1,2,figsize=(12,4))

    myerrorbar(toy[:,0],ax[0],bins=nbins,range=mrange)
    myerrorbar(toy[:,1],ax[1],bins=nbins,range=trange)

    if draw_pdf:
        m = np.linspace(*mrange,400)
        mN = (mrange[1]-mrange[0])/nbins

        bkgm = mpdf(m, *(ypars+mpars),comps=['bkg'])
        sigm = mpdf(m, *(ypars+mpars),comps=['sig'])
        totm = bkgm + sigm

        ax[0].plot(m, mN*bkgm, 'r--', label='Background')
        ax[0].plot(m, mN*sigm, 'g:' , label='Signal')
        ax[0].plot(m, mN*totm, 'b-' , label='Total PDF')

        t = np.linspace(*trange,400)
        tN = (trange[1]-trange[0])/nbins

        bkgt = tpdf(t, *(ypars+tpars),comps=['bkg'])
        sigt = tpdf(t, *(ypars+tpars),comps=['sig'])
        tott = bkgt + sigt

        ax[1].plot(t, tN*bkgt, 'r--', label='Background')
        ax[1].plot(t, tN*sigt, 'g:' , label='Signal')
        ax[1].plot(t, tN*tott, 'b-' , label='Total PDF')

    ax[0].set_xlabel('Invariant mass')
    ax[0].set_ylim(bottom=0)
    ax[0].legend()

    ax[1].set_xlabel('Decay time')
    ax[1].set_ylim(bottom=0)
    ax[1].legend()

    fig.tight_layout()

plot(toy)
_images/tutorial_5_0.png

Check the independence of our data

By computing the kendall rank coefficient and seeing how compatibile it is with 0

[4]:
kts = kendall_tau(toy[:,0],toy[:,1])
print('Kendall Tau:', pdg_format( kts[0], kts[1] ) )
plot_indep_scatter(toy[:,0],toy[:,1],reduction_factor=2);
Kendall Tau: -0.017 ± 0.010
_images/tutorial_7_1.png

Fit the toy in invariant mass

This provides us with estimates for the component shapes and the component yields

[5]:
# define mass pdf for iminuit fitting
def mpdf_min(x, Ns, Nb, mu, sg, lb):
    return (Ns+Nb, mpdf(x, Ns, Nb, mu, sg, lb) )

mi = Minuit( ExtendedUnbinnedNLL(toy[:,0], mpdf_min), Ns=Ns, Nb=Nb, mu=mu, sg=sg, lb=lb )
mi.limits['Ns'] = (0,Ns+Nb)
mi.limits['Nb'] = (0,Ns+Nb)
mi.limits['mu'] = mrange
mi.limits['sg'] = (0,mrange[1]-mrange[0])
mi.limits['lb'] = (0,10)

mi.migrad()
mi.hesse()
display(mi) # only valid for ipython notebooks
Migrad
FCN = -1.684e+05 Nfcn = 135
EDM = 7.37e-05 (Goal: 0.0002)
Valid Minimum No Parameters at limit
Below EDM threshold (goal x 10) Below call limit
Covariance Hesse ok Accurate Pos. def. Not forced
Name Value Hesse Error Minos Error- Minos Error+ Limit- Limit+ Fixed
0 Ns 5.11e3 0.11e3 0 1E+04
1 Nb 4.89e3 0.11e3 0 1E+04
2 mu 0.4981 0.0021 0 1
3 sg 0.1031 0.0022 0 1
4 lb 1.07 0.06 0 10
Ns Nb mu sg lb
Ns 1.22e+04 -7.12e+03 (-0.587) -0.00789 (-0.034) 0.124 (0.516) -1.44 (-0.203)
Nb -7.12e+03 (-0.587) 1.2e+04 0.00788 (0.034) -0.124 (-0.520) 1.44 (0.205)
mu -0.00789 (-0.034) 0.00788 (0.034) 4.46e-06 -2.09e-07 (-0.045) -3.35e-05 (-0.247)
sg 0.124 (0.516) -0.124 (-0.520) -2.09e-07 (-0.045) 4.72e-06 -2.39e-05 (-0.172)
lb -1.44 (-0.203) 1.44 (0.205) -3.35e-05 (-0.247) -2.39e-05 (-0.172) 0.00411

Construct the sweighter

Note that this will run much quicker if verbose=False and checks=False

[6]:
# define estimated functions
spdf = lambda m: mpdf(m,*mi.values,comps=['sig'])
bpdf = lambda m: mpdf(m,*mi.values,comps=['bkg'])

# make the sweighter
sweighter = SWeight( toy[:,0], [spdf,bpdf], [mi.values['Ns'],mi.values['Nb']], (mrange,), method='summation', compnames=('sig','bkg'), verbose=True, checks=True )
Initialising sweight with the summation method:
    PDF normalisations:
         0 5107.224408977729
         1 4892.156810093505
    W-matrix:
        [[1.38699027e-04 5.96207082e-05]
         [5.96207082e-05 1.42184207e-04]]
    A-matrix:
        [[ 8795.16287156 -3687.98933236]
         [-3687.98933236  8579.57828096]]
    Integral of w*pdf matrix (should be close to the
                identity):
        [[ 1.00011709e+00 -1.32635071e-04]
         [-1.57992838e-04  1.00004886e+00]]
    Check of weight sums (should match yields):
        Component  | sWeightSum |   Yield    |   Diff    |
        ---------------------------------------------------
          0        |  5107.2244 |  5107.2244 |    -0.00% |
          1        |  4892.1568 |  4892.1568 |     0.00% |

Construct the COW

Note that COW pdfs are always normalised in the COW numerator so that the W and A matrices tend to come out with elements of order one. This example uses a variance function of unity, I(m)=1, but also codes the case where I(m) = g(m), which is equiavalent to sweights.

[7]:
# unity
Im = 1
# sweight equiavlent
# Im = lambda m: mpdf(m,*mi.values) / (mi.values['Ns'] + mi.values['Nb'] )

# make the cow
cw = Cow(mrange, spdf, bpdf, Im, verbose=True)



Initialising COW:
    W-matrix:
        [[2.73492037 0.97050271]
          [0.97050271 1.07239023]]
    A-matrix:
        [[ 0.53861176 -0.48743839]
          [-0.48743839  1.37362336]]

Comparison of the sweight and COW methods

Compare the weight distributoins

[8]:
def plot_wts(x, sw, bw, title=None):
    fig,ax = plt.subplots()
    ax.plot(x, sw, 'b--', label='Signal')
    ax.plot(x, bw, 'r:' , label='Background')
    ax.plot(x, sw+bw, 'k-', label='Sum')
    ax.set_xlabel('Mass')
    ax.set_ylabel('Weight')
    if title: ax.set_title(title)
    fig.tight_layout()

for meth, cls in zip( ['SW','COW'], [sweighter,cw] ):

    # plot weights
    x = np.linspace(*mrange,400)
    swp = cls.get_weight(0,x)
    bwp = cls.get_weight(1,x)
    plot_wts(x, swp, bwp, meth)
_images/tutorial_16_0.png
_images/tutorial_16_1.png

Fit the weighted data in decay time and correct the covariance

[9]:
# define weighted nll
def wnll(tlb, tdata, wts):
    sig  = expon(trange[0],tlb)
    sigN = np.diff( sig.cdf(trange) )
    return -np.sum( wts * np.log( sig.pdf( tdata ) / sigN ) )

# define signal only time pdf for cov corrector
def tpdf_cor(x, tlb):
    return tpdf(x,1,0,tlb,['sig'])

flbs=[]
for meth, cls in zip( ['SW','COW'], [sweighter,cw] ):

    print('Method:', meth)
    # get the weights
    wts = cls.get_weight(0,toy[:,0])

    # define the nll
    nll = lambda tlb: wnll(tlb, toy[:,1], wts)

    # do the minimisation
    tmi = Minuit( nll, tlb=tlb )
    tmi.limits['tlb'] = (1,3)
    tmi.errordef = Minuit.LIKELIHOOD
    tmi.migrad()
    tmi.hesse()

    # and do the correction
    fval = np.array(tmi.values)
    flbs.append(fval[0])
    fcov = np.array( tmi.covariance.tolist() )

    # first order correction
    ncov = approx_cov_correct(tpdf_cor, toy[:,1], wts, fval, fcov, verbose=False)

    # second order correction
    hs  = tpdf_cor
    ws  = lambda m: cls.get_weight(0,m)
    W   = cls.Wkl

    # these derivatives can be done numerically but for the sweights / COW case it's straightfoward to compute them
    ws = lambda Wss, Wsb, Wbb, gs, gb: (Wbb*gs - Wsb*gb) / ((Wbb-Wsb)*gs + (Wss-Wsb)*gb)
    dws_Wss = lambda Wss, Wsb, Wbb, gs, gb: gb * ( Wsb*gb - Wbb*gs ) / (-Wss*gb + Wsb*gs + Wsb*gb - Wbb*gs)**2
    dws_Wsb = lambda Wss, Wsb, Wbb, gs, gb: ( Wbb*gs**2 - Wss*gb**2 ) / (Wss*gb - Wsb*gs - Wsb*gb + Wbb*gs)**2
    dws_Wbb = lambda Wss, Wsb, Wbb, gs, gb: gs * ( Wss*gb - Wsb*gs ) / (-Wss*gb + Wsb*gs + Wsb*gb - Wbb*gs)**2

    tcov = cov_correct(hs, [spdf,bpdf], toy[:,1], toy[:,0], wts, [mi.values['Ns'],mi.values['Nb']], fval, fcov, [dws_Wss,dws_Wsb,dws_Wbb],[W[0,0],W[0,1],W[1,1]], verbose=False)
    print('Method:', meth, f'- covariance corrected {fval[0]:.2f} +/- {fcov[0,0]**0.5:.2f} ---> {fval[0]:.2f} +/- {tcov[0,0]**0.5:.2f}')



Method: SW
Method: SW - covariance corrected 1.71 +/- 0.14 ---> 1.71 +/- 0.19
Method: COW
Method: COW - covariance corrected 1.64 +/- 0.13 ---> 1.64 +/- 0.18
[10]:
### Plot the weighted decay distributions and the fit result
def plot_tweighted(x, wts, wtnames=[], funcs=[]):

    fig, ax = plt.subplots()

    t = np.linspace(*trange,400)
    N = (trange[1]-trange[0])/50

    for i, wt in enumerate(wts):
        label = None
        if i<len(wtnames): label = wtnames[i]
        myerrorbar(x, ax, bins=50, range=trange, wts=wt, label=label, col=f'C{i}')
        if i<len(funcs):
            ax.plot(t,N*funcs[i](t),f'C{i}-')

    ax.legend()
    ax.set_xlabel('Time')
    ax.set_ylabel('Weighted Events')
    fig.tight_layout()

swf  = lambda t: tpdf(t, mi.values['Ns'], 0, flbs[0], comps=['sig'] )
cowf = lambda t: tpdf(t, mi.values['Ns'], 0, flbs[1], comps=['sig'] )
sws  = sweighter.get_weight(0, toy[:,0])
scow = cw.get_weight(0, toy[:,0])

plot_tweighted(toy[:,1], [sws,scow], ['SW','COW'], funcs=[swf,cowf] )


_images/tutorial_19_0.png

Write the weights back into a file

For example into a pandas dataframe or root file

[11]:
import pandas as pd

df = pd.DataFrame( )
df['mass'] = toy[:,0]
df['time'] = toy[:,1]
df['sw_sws'] = sweighter.get_weight(0, df['mass'].to_numpy() )
df['sw_bws'] = sweighter.get_weight(1, df['mass'].to_numpy() )
df['cow_sws'] = cw.get_weight(0, df['mass'].to_numpy() )
df['cow_bws'] = cw.get_weight(1, df['mass'].to_numpy() )

import uproot

with uproot.recreate('outf.root') as f:
    f['tree'] = df

print(df)

          mass      time    sw_sws    sw_bws   cow_sws   cow_bws
0     0.276401  0.287028 -0.129313  1.129223 -0.372759  1.446466
1     0.302426  0.426394  0.150924  0.848999 -0.220921  1.282285
2     0.578586  0.449921  1.149476 -0.149511  1.099331 -0.159718
3     0.373953  0.741222  0.839109  0.160843  0.481153  0.576620
4     0.515549  0.598659  1.244910 -0.244940  1.590285 -0.553110
...        ...       ...       ...       ...       ...       ...
9995  0.846222  0.426296 -0.706615  1.706501 -0.332575  0.950587
9996  0.166248  0.153495 -0.711678  1.711563 -0.631136  1.801133
9997  0.535665  0.266288  1.231984 -0.232015  1.494705 -0.483188
9998  0.357202  0.365763  0.708016  0.291931  0.282588  0.772361
9999  0.643105  0.593656  0.832079  0.167873  0.364086  0.456587

[10000 rows x 6 columns]