383 lines
14 KiB
Python
Executable File
383 lines
14 KiB
Python
Executable File
import warnings
|
|
from math import pi
|
|
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
import zfit
|
|
from zfit import z
|
|
|
|
import yaml
|
|
import sympy
|
|
|
|
from b2kstll.models.Kpi import KpiPWavePDF, KpiSWavePDF, ReKpiPandSWavePDF, ImKpiPandSWavePDF
|
|
from b2kstll.models.acceptance import AccPDF
|
|
from b2kstll.utils import get_coeffs
|
|
|
|
import analysis.efficiency
|
|
from analysis.efficiency import load_efficiency_model
|
|
|
|
from particle import Particle
|
|
|
|
from b2kstll.models.angular_sympy import angular_func_sympy
|
|
|
|
|
|
class AngularPDF(zfit.pdf.BasePDF):
|
|
""" Generic angular PDF built from input function and corresponding parameters
|
|
|
|
Args:
|
|
ArbitraryFNC (function: Any arbitrary angular function
|
|
params (params): parameters of interest in the given function
|
|
obs (`zfit.Space`):
|
|
name (str):
|
|
dtype (tf.DType):
|
|
"""
|
|
def __init__(self, obs, params, func, name):
|
|
super().__init__(obs=obs, params=params, name=name+"PDF")
|
|
self._user_func = func
|
|
|
|
_N_OBS = 3
|
|
VALID_LIMITS = {'costheta_l': (-1, 1),
|
|
'costheta_k': (-1, 1),
|
|
'phi': (-pi, pi)}
|
|
|
|
def _pdf(self, x, norm_range):
|
|
costheta_l, costheta_k, phi = z.unstack_x(x)
|
|
|
|
pdf = self._user_func(**self.params, costheta_l=costheta_l, costheta_k=costheta_k, phi=phi)
|
|
return pdf
|
|
|
|
# Calculate Normalisation
|
|
def get_normalisation(FL, upper_ctk):
|
|
kwargs = dict(
|
|
FL = sympy.Symbol('FL'),
|
|
AFB = sympy.Symbol('AFB'),
|
|
S3 = sympy.Symbol('S3'),
|
|
S4 = sympy.Symbol('S4'),
|
|
S5 = sympy.Symbol('S5'),
|
|
S7 = sympy.Symbol('S7'),
|
|
S8 = sympy.Symbol('S8'),
|
|
S9 = sympy.Symbol('S9'),
|
|
costheta_l = sympy.Symbol('costheta_l'),
|
|
costheta_k = sympy.Symbol('costheta_k'),
|
|
phi = sympy.Symbol('phi')
|
|
)
|
|
|
|
sympy_val = angular_func_sympy(**kwargs)
|
|
integral = sympy.integrate(sympy_val,
|
|
(kwargs['costheta_l'], -1, 1),
|
|
(kwargs['costheta_k'], -1, upper_ctk),
|
|
(kwargs['phi'], -sympy.pi, sympy.pi),)
|
|
|
|
fun = sympy.lambdify((kwargs["FL"]),integral, 'tensorflow')
|
|
|
|
return fun(FL)
|
|
|
|
UPPER_CTK = 0.85
|
|
|
|
def PWave(FL, S3, S4, S5, AFB, S7, S8, S9, costheta_l, costheta_k, phi, backend=tf):
|
|
|
|
"""Full d4Gamma/dq2dOmega for Bd -> Kst ll (l=e,mu)
|
|
|
|
Angular distribution obtained in the total PDF (using LHCb convention JHEP 02 (2016) 104)
|
|
i.e. the valid of the angles is given for
|
|
- phi: [-pi, pi]
|
|
- theta_K: [-pi, pi]
|
|
- theta_l: [-pi, pi]
|
|
|
|
The function is normalized over a finite range and therefore a PDF.
|
|
|
|
Args:
|
|
FL (`zfit.Parameter`): Fraction of longitudinal polarisation of the Kst
|
|
S3 (`zfit.Parameter`): A_perp^2 - A_para^2 / A_zero^2 + A_para^2 + A_perp^2 (L, R)
|
|
S4 (`zfit.Parameter`): RE(A_zero*^2 * A_para^2) / A_zero^2 + A_para^2 + A_perp^2 (L, R)
|
|
S5 (`zfit.Parameter`): RE(A_zero*^2 * A_perp^2) / A_zero^2 + A_para^2 + A_perp^2 (L, R)
|
|
AFB (`zfit.Parameter`): Forward-backward asymmetry of the di-lepton system (also i.e. 3/4 * S6s)
|
|
S7 (`zfit.Parameter`): IM(A_zero*^2 * A_para^2) / A_zero^2 + A_para^2 + A_perp^2 (L, R)
|
|
S8 (`zfit.Parameter`): IM(A_zero*^2 * A_perp^2) / A_zero^2 + A_para^2 + A_perp^2 (L, R)
|
|
S9 (`zfit.Parameter`): IM(A_perp*^2 * A_para^2) / A_zero^2 + A_para^2 + A_perp^2 (L, R)
|
|
obs (`zfit.Space`):
|
|
name (str):
|
|
dtype (tf.DType):
|
|
"""
|
|
|
|
sintheta_k = backend.sqrt(1.0 - costheta_k * costheta_k)
|
|
sintheta_l = backend.sqrt(1.0 - costheta_l * costheta_l)
|
|
|
|
sintheta_2k = (1.0 - costheta_k * costheta_k)
|
|
sintheta_2l = (1.0 - costheta_l * costheta_l)
|
|
|
|
sin2theta_k = (2.0 * sintheta_k * costheta_k)
|
|
cos2theta_l = (2.0 * costheta_l * costheta_l - 1.0)
|
|
sin2theta_l = (2.0 * sintheta_l * costheta_l)
|
|
|
|
pdf = ((3.0 / 4.0) * (1.0 - FL) * sintheta_2k +
|
|
FL * costheta_k * costheta_k +
|
|
(1.0 / 4.0) * (1.0 - FL) * sintheta_2k * cos2theta_l +
|
|
-1.0 * FL * costheta_k * costheta_k * cos2theta_l +
|
|
S3 * sintheta_2k * sintheta_2l * backend.cos(2.0 * phi) +
|
|
S4 * sin2theta_k * sin2theta_l * backend.cos(phi) +
|
|
S5 * sin2theta_k * sintheta_l * backend.cos(phi) +
|
|
(4.0 / 3.0) * AFB * sintheta_2k * costheta_l +
|
|
S7 * sin2theta_k * sintheta_l * backend.sin(phi) +
|
|
S8 * sin2theta_k * sin2theta_l * backend.sin(phi) +
|
|
S9 * sintheta_2k * sintheta_2l * backend.sin(2.0 * phi))
|
|
|
|
#return (9.0 / (32.0 * pi)) * pdf
|
|
return pdf / get_normalisation(FL = FL, upper_ctk = UPPER_CTK)
|
|
|
|
|
|
def SWave(costheta_l, costheta_k, phi, backend=tf):
|
|
|
|
"""Full d4Gamma/dq2dOmega for Swave part only
|
|
|
|
Angular distribution obtained in the total PDF (using LHCb convention JHEP 02 (2016) 104)
|
|
i.e. the valid of the angles is given for
|
|
- phi: [-pi, pi]
|
|
- theta_K: [-pi, pi]
|
|
- theta_l: [-pi, pi]
|
|
|
|
The function is normalized over a finite range and therefore a PDF.
|
|
|
|
Args:
|
|
obs (`zfit.Space`):
|
|
name (str):
|
|
dtype (tf.DType):
|
|
"""
|
|
|
|
sintheta_2l = (1.0 - costheta_l * costheta_l)
|
|
pdf = (3.0 / (16.0 * pi)) * sintheta_2l
|
|
|
|
return pdf
|
|
|
|
|
|
def PandSWaveInterference(Ss1, Ss2, Ss3, Ss4, Ss5, costheta_l, costheta_k, phi, backend=tf):
|
|
|
|
"""Full d4Gamma/dq2dOmega for the Real interference term between P and S waves
|
|
|
|
Angular distribution obtained in the total PDF (using LHCb convention JHEP 02 (2016) 104)
|
|
i.e. the valid of the angles is given for
|
|
- phi: [-pi, pi]
|
|
- theta_K: [-pi, pi]
|
|
- theta_l: [-pi, pi]
|
|
|
|
The function is normalized over a finite range and therefore a PDF.
|
|
|
|
Args:
|
|
Ss1, Ss2, Ss3, Ss4, Ss5 (`zfit.Parameter`)
|
|
obs (`zfit.Space`):
|
|
name (str):
|
|
dtype (tf.DType):
|
|
"""
|
|
|
|
sintheta_l = backend.sqrt(1.0 - costheta_l * costheta_l)
|
|
sintheta_k = backend.sqrt(1.0 - costheta_k * costheta_k)
|
|
|
|
sintheta_2l = (1.0 - costheta_l * costheta_l)
|
|
sin2theta_l = (2.0 * sintheta_l * costheta_l)
|
|
|
|
pdf = (Ss1 * sintheta_2l * costheta_k +
|
|
Ss2 * sin2theta_l * sintheta_k * backend.cos(phi) +
|
|
Ss3 * sintheta_l * sintheta_k * backend.cos(phi) +
|
|
Ss4 * sintheta_l * sintheta_k * backend.sin(phi) +
|
|
Ss5 * sin2theta_l * sintheta_k * backend.sin(phi))
|
|
|
|
return (3.0 / (16.0 * pi)) * pdf
|
|
|
|
|
|
class AngularPandSWavePDF(zfit.pdf.BaseFunctor):
|
|
"""Full d4Gamma/dq2dOmega for both spin 0 and 1 for Bd -> Kst ll (l=e,mu) decays
|
|
|
|
Angular distribution obtained in the total PDF (using LHCb convention JHEP 02 (2016) 104)
|
|
i.e. the valid of the angles is given for
|
|
- phi: [-pi, pi]
|
|
- theta_K: [0, pi]
|
|
- theta_l: [0, pi]
|
|
|
|
The function is normalized over a finite range and therefore a PDF.
|
|
|
|
Args:
|
|
AngularPWavePDF (`zfit.pdf.ZPDF`: P-wave normalised PDF
|
|
AngularSWavePDF (`zfit.pdf.ZPDF`: S-wave interference part normalised PDF
|
|
FS (`zfit.Parameter`): S-wave fraction, i.e. |A^L_S|^2 + |A^R_S|^2/ Sum(|A^{S,R}_{S,zero,para,perp}|^2)
|
|
obs (`zfit.Space`):
|
|
name (str):
|
|
dtype (tf.DType):
|
|
"""
|
|
|
|
def __init__(self, PWave, SWave, PandSWaveInterference, FS, obs):
|
|
super().__init__(pdfs=[PWave, SWave, PandSWaveInterference], name="AngularPandSWavePDF", params={'FS': FS}, obs=obs)
|
|
|
|
_N_OBS = 3
|
|
VALID_LIMITS = {'costheta_l': (-1, 1),
|
|
'costheta_k': (-1, 1),
|
|
'phi': (-pi, pi)}
|
|
|
|
def _pdf(self, x, norm_range):
|
|
FS = self.params['FS']
|
|
|
|
PWave, SWave, PandSWaveInterference = self.pdfs
|
|
|
|
pdf = (1 - FS) * PWave.pdf(x) + FS * SWave.pdf(x) + PandSWaveInterference.pdf(x)
|
|
|
|
return pdf
|
|
|
|
|
|
|
|
# A bit of handling
|
|
class B2Kstll:
|
|
"""Factory class for B->K*ll angular fit"""
|
|
|
|
PDFS = {'PWave': (PWave, ['FL', 'S3', 'S4', 'S5', 'AFB', 'S7', 'S8', 'S9']),
|
|
'SWave': (SWave, []),
|
|
'PandSWaveInterference' : (PandSWaveInterference, ['Ss1', 'Ss2', 'Ss3', 'Ss4', 'Ss5']),
|
|
#'PandSWaveInterference5D': (PandSWaveInterference5D, ['Ss1re', 'Ss2re', 'Ss3re', 'Ss4re', 'Ss5re',
|
|
# 'Ss1im', 'Ss2im', 'Ss3im', 'Ss4im', 'Ss5im']),
|
|
'PandSWave' : (AngularPandSWavePDF, ['FS'])}
|
|
#'PandSWave5D': (AngularPandSWavePDF5D, ['FS'])}
|
|
|
|
KPI = {'PWave': (KpiPWavePDF, ['Kst_892'], [313]),
|
|
'SWave': (KpiSWavePDF, ['Kst_1430'], [10311]),
|
|
'RePandSWave': (ReKpiPandSWavePDF, ['Kst_892', 'Kst_1430'], [313, 10311]),
|
|
'ImPandSWave': (ImKpiPandSWavePDF, ['Kst_892', 'Kst_1430'], [313, 10311])}
|
|
|
|
|
|
def __init__(self, cosThetaL, cosThetaK, Phi, q2_min=None, q2_max=None):
|
|
"""Initialize observable names."""
|
|
self._obs_names = {'costheta_l': cosThetaL,
|
|
'costheta_k': cosThetaK,
|
|
'phi': Phi}
|
|
|
|
costheta_l = zfit.Space(cosThetaL, limits=(-1.0, 1.0))
|
|
costheta_k = zfit.Space(cosThetaK, limits=(-1.0, UPPER_CTK))
|
|
phi = zfit.Space(Phi, limits=(-pi, pi))
|
|
self._spaces = costheta_l * costheta_k * phi
|
|
|
|
self.params = {}
|
|
|
|
self.acc = None
|
|
self.acc_integral = None
|
|
self.Kpi = None
|
|
self.q2_bin_acc = None
|
|
self.q2_min = q2_min
|
|
self.q2_max = q2_max
|
|
|
|
def set_q2_bin_acceptance(self, q2_bin_acc):
|
|
"""Set the q2 value used to slice the acceptance function for parametrization in 4D"""
|
|
self.q2_bin_acc = q2_bin_acc
|
|
|
|
def set_acc(self, file_name, q2_bin_acc=None):
|
|
"""Set an acceptance function to the signal PDF."""
|
|
|
|
# Load coefficients from MoM
|
|
#coeffs_yaml = load_efficiency_model(file_name)
|
|
#coeffs = coeffs_yaml.get_coefficients()
|
|
|
|
self.acc,self.accmodel = get_coeffs(file_name)
|
|
self.set_q2_bin_acceptance(q2_bin_acc)
|
|
|
|
|
|
def set_acc_integral(self, file_name, str_order):
|
|
"""Set the integral of the PDF x acceptance."""
|
|
|
|
# Load integral from a yaml file
|
|
with open(file_name, 'r') as f:
|
|
dict_int = yaml.load(f, Loader=yaml.FullLoader)
|
|
|
|
integral = dict_int[0][str_order]
|
|
self.acc_integral = sympy.sympify(integral)
|
|
|
|
|
|
def set_Kpi(self, name):
|
|
""" Set the Kpi invariant mass system."""
|
|
|
|
#self._has_Kpi = True
|
|
self.Kpi = name
|
|
|
|
|
|
def get_pdf(self, name):
|
|
"""Get corresponding PDF and paramaters."""
|
|
fnc_class, param_names = self.PDFS[name]
|
|
|
|
def get_params(param_list):
|
|
out = {}
|
|
for param in param_list:
|
|
if param not in self.params:
|
|
if param == 'FL':
|
|
config = [0.6, 0, 1]
|
|
elif param == 'FS':
|
|
config = [0.1, 0, 1]
|
|
else:
|
|
config = [0.0, -2, 2]
|
|
self.params.update({param: zfit.Parameter(param, *config)})
|
|
out[param] = self.params[param]
|
|
return out
|
|
|
|
obs = self._spaces
|
|
|
|
def get_Kpi_pdf(name, model):
|
|
"""Set the corresponding PDF for the Kpi system and parameters."""
|
|
pdf_class, param_names, param_id = self.KPI[model]
|
|
|
|
def get_res(param_list, param_id):
|
|
out = {}
|
|
for param, idparam in zip(param_names, param_id):
|
|
if not any([i for i in self.params if param in i.lower()]):
|
|
m = Particle.from_pdgid(idparam).mass * 1e-03
|
|
w = Particle.from_pdgid(idparam).width * 1e-03
|
|
self.params.update({param+"_M": zfit.Parameter(param+"_M", m, floating=False)})
|
|
self.params.update({param+"_W": zfit.Parameter(param+"_W", w, floating=False)})
|
|
out[param+"_M"] = self.params[param+"_M"]
|
|
out[param+"_W"] = self.params[param+"_W"]
|
|
return out
|
|
|
|
# Add Kpi dimension
|
|
mKpi = zfit.Space(name, limits=(0.79555, 0.99555))
|
|
self._spaces *= mKpi
|
|
|
|
params = get_res(param_names, param_id)
|
|
pdf = pdf_class(obs=mKpi, params=params, name=model, q2_bin=(self.q2_min+self.q2_max)/2.)
|
|
return pdf
|
|
|
|
|
|
if name == 'PandSWave':
|
|
pdf_tot, params_tot = self.PDFS[name]
|
|
get_params(params_tot)
|
|
|
|
listPDF = dict.fromkeys(['PWave', 'SWave', 'PandSWaveInterference'])
|
|
for iPDF in listPDF:
|
|
ifncWave, iparamWave = self.PDFS[iPDF]
|
|
initParams = get_params(iparamWave)
|
|
listPDF[iPDF] = AngularPDF(obs=obs, params=initParams, func=ifncWave, name=iPDF)
|
|
|
|
if self.Kpi:
|
|
pdf_Kpi = get_Kpi_pdf(self.Kpi, iPDF)
|
|
listPDF[iPDF] *= pdf_Kpi
|
|
|
|
params = self.params
|
|
pdf = pdf_tot(obs=obs, PWave=listPDF["PWave"], SWave=listPDF["SWave"], PandSWaveInterference=listPDF["PandSWaveInterference"], FS=self.params['FS'])
|
|
if self.acc is not None:
|
|
pdf = AccPDF(pdf, self.acc, self.accmodel, obs=obs, dict_names=self._obs_names, q2_slice=self.q2_bin_acc, name=name)
|
|
|
|
else:
|
|
params = get_params(param_names)
|
|
pdf = AngularPDF(obs=obs, params=params, func=fnc_class, name=name)
|
|
|
|
if self.acc is not None:
|
|
pdf = AccPDF(pdf, self.acc, self.accmodel, obs=obs, dict_names=self._obs_names, q2_slice=self.q2_bin_acc, name=name)
|
|
|
|
if self.acc_integral:
|
|
|
|
tf_expr = sympy.lambdify(self.acc_integral.free_symbols, self.acc_integral, 'tensorflow')
|
|
integral = integral_acceptance(params, tf_expr)
|
|
pdf.register_analytic_integral(func=integral, limits=zfit.Space(axes=(0,1,2), limits=obs.limits))
|
|
|
|
if self.Kpi:
|
|
pdf_Kpi = get_Kpi_pdf(self.Kpi, name)
|
|
pdf = pdf * pdf_Kpi
|
|
|
|
params = self.params
|
|
obs = self._spaces
|
|
|
|
return obs, pdf, params
|
|
|