\section{Beyond standard fitting}
\label{sec:beyond standard fitting}

In the previous Sections, we have seen the capabilities of \zfit{} for fitting 
a model to data. The inherent flexibility and the powerful functionality of the 
core parts allow for the
library to be extended beyond the usual feature set of general model fitting 
libraries in HEP.
To perform more specialized fits, the components of \zfit{} can be 
used as building blocks of a higher level fitter. In order to keep the size and 
features of \zfit{} 
at a reasonable level, the whole project 
is split into multiple packages that are tightly coupled to each other. In this 
Section, we will have a closer look at the different packages that extend 
\zfit{} beyond the simple model fitting we have seen so far.

The main package is the core library \zfit{}, as described in the thesis up to 
now, that provides all the fundamental 
building blocks and is in itself a self-contained fitting library. 
Content-wise, 
\zfit{} offers a more field agnostic selection and does not contain models and 
tools too specific for HEP. The focus is on a stable and solid 
implementation together with the API and format definitions.

More field specific content can be added in other libraries such as \zphys{}. 
This repository is meant to be the place for community 
contributions focusing on HEP-specific content. Guidelines, examples and
automated tests are planned to be in place in order to lower 
the threshold 
for contributions. Furthermore, it will also contain 
functions dealing with physical quantities such as kinematics, which can serve 
as 
building blocks for models. Higher level interfaces that use \zfit{} to build 
specific models are also intended to be placed in this repositories, for 
example a whole amplitude analysis framework currently under development which
will be explained in the following Section.

Additionally, the project also contains libraries that \zfit{} depends on but 
which themselves are self-contained and which are factored out of the 
core- and extension libraries to standalone packages. This allows other 
projects than \zfit{} to make use of them as well. An example is the 
implementation of 
\texttt{phasespace}, which generates four momenta of particles from a 
decay taking 
into account the correct kinematics. The package will be explained in more 
detail in Sec. \ref{sec:phasespace}.


\subsection{Amplitude fits}
\label{sec:dalitz}

\zamp{} is a higher level fitting library, which is currently under 
active development and in its early stage, built with elements from \zfit{} in 
order to perform amplitude analyses, including Dalitz. 
In the following, the extension within the scope 
of the \zfit{} project is discussed by showing an implementation example of the 
decay \decay{\Dz}{\Kp\pim\piz}. As part of this, an additional library will be 
introduced, \zphasespace, which covers an essential part of being able to 
generate amplitude fits.

In order to perform amplitude fits, \zamp{} contains
\begin{itemize}
	\item shapes such as the Breit-Wigner distribution that can be used to 
	model resonances;
	\item helper classes to build PDFs such as some that efficiently cache 
	intermediate results specific to amplitude fits, and
	\item a higher level interface to build amplitude fits in a transparent 
	manner, similar to other specialised amplitude fitters in HEP. This removes 
	the need of low level handling \textit{but} keeps the flexibility to 
	replace any part of it by a custom object.
\end{itemize}

In order to perform toy studies as described in \ref{sec:performance}, an 
efficient way of sampling from a model is needed. While accept-reject is an 
universal and good working way, the efficiency can be very low in higher 
dimensional spaces and/or with peaky model shapes. To avoid inefficiencies, 
samples can be produced by using importance sampling as described in Appendix 
\ref{appendix:sampling techniques}, which requires a distribution 
approximately 
resembling the shape of the model. In \zamp{} the approach is to use the 
kinematic constituents of the decay particles. They are then transformed to the 
desired variables that are used in the model. This is a very general approach 
that allows all sorts of variables to be constructed from the kinematics. On 
the downside, in order to produce the particles with realistic kinematics, a 
phasespace generator is needed. While there is an implementation in \root 
called the TGenPhasespace class, no equivalent is available in pure Python. 
Therefore, a library porting the above with an extension to real world 
experiment kinematics and implemented using TF has been created.

More details on the \zamp{} library, especially on the higher level interface,
will be provided in a future paper and goes beyond the scope of this thesis.

\subsection{phasespace}
\label{sec:phasespace}

The kinematics of a particle are a four-dimensional tuple representing its 
four-momentum. In a decay of a parent particle to lighter children particles, 
their kinematics are constrained by the fundamental physical laws of momentum 
and energy conservation.
With
\begin{align*}
	p^{0,1,2} = p^{x, y, z} ~~~~ p^3 = \sqrt{p^2 + m^2c^2}
\end{align*}
this reads as
\begin{equation}
\label{eq:momentum_conservation}
	p_{parent}^{\mu} = \sum_i p_{i}^{\mu}
\end{equation}
with $p_i$ being the four-momenta of all children particles. For a two-body 
decay $A\rightarrow BC$, there are six free variables from the momenta of the 
two children. Using Eq. \ref{eq:momentum_conservation} there are two degrees of 
freedom left in the decay kinematics that lead to a distribution. Furthermore, 
the mass is only fixed for stable particles but is a distribution for 
resonances. To sample from the phasespace within the \zfit{} project 
a 
package named \zphasespace{} was created. The purpose of it is to generate 
arbitrary decays obeying the physical constraints discussed before and 
expressed in Eq. \ref{eq:momentum_conservation}. Since it uses TF as the 
backend, it integrates seamlessly into \zfit{}.

The algorithm used in \zphasespace{} is the Raubold and Lynch method for n-body 
events as described in \cite{James:1968gu}. In principle, the algorithm builds 
a tree where every node has two leaves representing a two-body decay. It 
calculates the available kinematic energy that will be assigned to the 
particles from the difference of the parent rest mass and the sum of all the 
children rest masses. The latter are either fixed or drawn from the mass 
distribution in case of short-lived particles

\begin{equation*}
	E_{kin} = m_{parent} - \sum_{i} f_{i}^{mass} (m_{i}^{min}, m_{i}^{max})
\end{equation*}
with $f_{i}^{mass}$ being the mass distribution of the particle $i$, 
$m_i^{min}$ 
the minimum mass recursively determined from the children masses of 
particle i and $m_i^{max}$ the available energy from the top particle minus the 
other parent particles of particle $i$. The remaining $E_{kin}$ is randomly 
split 
into fractions along all the decaying particles in the tree where each fraction 
is the kinetic energy for the boost of this particle.

Given the mass of each particle in the tree and its kinetic energy, particles 
are recursively generated starting from the top of the tree, \emph{i.e.}, with 
the 
lightest particles. Each parent 
particle randomly generates the two childrens in the available phasespace. Then 
the whole decay tree is boosted to the parent momentum. This continues until 
the top particle is reached, which is not boosted by default. However, an 
additional argument to the generation method allows to also boost it. This can 
be used to reproduce the physics happening in actual colliders.

\subsubsection*{Usage}
The library has a main class \pyth{GenParticle} that describes the particles 
mass, 
either fixed or as a distribution, and has a method to set the children 
particle it decays to. This allows to build an arbitrary decay chain in an 
object-oriented way. As an example, the decay of \decay{\Dz}{\Kstar(892) 
(\rightarrow{\Kp\pim})\piz}, which will be implemented as an amplitude fit in 
the next Section, would be implemented as following

\begin{center}
\begin{minipage}{\textwidth}
\label{py:reso_particle}
\begin{python}
import phasespace as phsp

kplus = phsp.GenParticle("K+", mass=KPLUS_MASS)
piplus = phsp.GenParticle("Pi+", mass=PIPLUS_MASS)
piminus = phsp.GenParticle("Pi-", mass=PIMINUS_MASS)
kstar = phsp.GenParticle("K*", mass=kstar_mass_func)
dzero = phsp.GenParticle("D0", mass=DZERO_MASS)

kstar.set_children(kplus, piminus)
dzero.set_children(kstar, piplus)
\end{python}
\end{minipage}
\end{center}

Here, \pyth{kstar_mass_func} is a function sampling from the mass distribution of a \Kstar. This is a custom function that can be implemented by the user.

Having specified the decay chain, two methods can be used to generate the 
actual particles: Either \pyth{generate_tensor} which returns the four momenta 
as a Tensor and can be used directly inside \zfit{}, or alternatively, 
\pyth{generate} returns the same information but as a numpy ndarray. 
Internally, a TF session is called and the computation is run.

\begin{center}
\begin{minipage}{\textwidth}
\begin{python}
weights, particles = dzero.generate(1000)  # generate 1000 particles
\end{python}
\end{minipage}
\end{center}

The weights correspond to every single event and quantify the probability of 
the generated event. The returned \pyth{particles} is a 
dictionary containing the 
momentum of each particle. For example

\begin{center}
\begin{minipage}{\textwidth}
\begin{python}
kstar_kinematics = particles['K*']
kstar_x = kstar_kinematics[:, 0]
kstar_mass = kstar_kinematics[:, 3]
\end{python}
\end{minipage}
\end{center}
and the format of the kinematic is (number of events, components).

As a shortcut for simple decays, a high level function \pyth{generate_decay} is 
available which allows to describe a decay with stacked lists of masses.

\subsection{Dalitz implementation}

As an example of an amplitude fit, a Dalitz analysis of the decay 
\decay{\Dz}{\Kp\pim\piz} is implemented using some parts of \zamp{} together 
with \zfit{} and \zphasespace. This 
implementation of the resonances and their shapes is done following the 
analysis in Ref. \cite{PhysRevLett.103.211801}.
The amplitude is built using the isobar approach, which calculates the coherent 
sum of the 
individual contributions, the intermediate resonances. The implemented 
resonances are $\rhomeson (770)$, $\rhomeson (1700)$, $\Kstarz (892) $, 
$\Kstarp (892)$, $\Kstarz (1430) $, $\Kstarp (1430)$ and
$K_2^*(1430)$

For amplitude analyses, dedicated fitting libraries exist and using a general 
purpose fitter like \zfit{} for this case is rather special. In the following, 
we will go through the elements that are needed from a top-down approach and 
see how the problem can be separated into smaller pieces using the 
functionality of \zfit{}. Although unimportant technicalities are left away and 
the example is a sketch only of the actual implementation, 
it is going to be rather advanced and involves functionality only 
explained in Appendices.

The observables in this case are the invariant masses of pairs of the different 
children 
particles, which are $m_{\Kp\pim}$, $m_{\Kp\piz}$ and $m_{\pim\piz}$. The whole 
PDF $f_{tot}$ is of the form
\begin{equation}
\label{eq:total amp}
	f_{tot}(m_{\Kp\pim}, m_{\Kp\piz}, m_{\pim\piz}) = \abs{\sum_{i}^{n_{amp}} 
	c_i A_i(m_{\Kp\pim}, m_{\Kp\piz}, m_{\pim\piz})}^2
\end{equation}
with $c_i$ being the complex coefficient of each amplitude $A_i$. Expanding 
this term contains cross ($i\neq j$) and square ($i=j$) terms $b_{ij}$ of the 
form
\begin{equation*}
	b_{ij} = c_i c_j^* A_i A_j^*
\end{equation*}

To implement this in \zfit{}, a custom class is created that takes the 
coefficients, 
amplitudes and a \pyth{Func} called \pyth{AmplitudeProduct} which calculates 
the above products $b_{ij}$. Before we will look at these three components, 
we can build the global PDF as defined in Eq. \ref{eq:total amp}

\begin{center}
\begin{minipage}{\textwidth}
\begin{python}
class SumAmplitudeSquaredPDF(zfit.pdf.BaseFunctor):
	def __init__(self, obs, coefs, amps,...)
		self._cross_terms = [AmplitudeProduct(c1, c2, a1, a2)
							for (c1, a1), (c2, a2) 
							in combinations(coefs, amps)]
		
		self._square_terms = [AmplitudeProduct(coef, coef, amp, amp)
				     		for coef, amp in zip(coefs, amps)]
		...
				     		
	def _unnormalized_pdf(self, x):
		value = tf.reduce_sum([2. * amp.func(x=x) for amp in self._cross_terms] 
					    	  +[amp.func(x=x) for amp in self._squared_terms],
					    	  axis=0)  # over which axis to sum
		return tf.real(value)  # convert it to a real number
\end{python}
\end{minipage}
\end{center}

As the next piece, we need to have the amplitude product. Since this is complex 
in general, it is again split into smaller parts that represent the real and 
imaginary parts by a class \pyth{AmplitudeProductProjection}.

\begin{center}
\begin{minipage}{\textwidth}
\begin{python}
class AmplitudeProduct(zfit.func.BaseFunctorFunc):
	def __init__(self, coef1, coef2, amp1, amp2,...):
		prod_real = AmplitudeProductProjection(amp1, amp2, proj=tf.math.real)
		prod_img = AmplitudeProductProjection(amp1, amp2, proj=tf.math.imag)
		super().__init__(funcs=[prod_real, prod_imag], 
						 params={'coef1': coef1, 'coef2', coef2}, ...)
		...
		
	def _func(self, x)
		coef1 = self.params['coef1']
		coef2 = self.params['coef2']
  	    prod_real, prod_imag = self.funcs	
  	    coeffs = coef1 * tf.conj(coef2)
		return  coeffs * tf.complex(prod_real.func(x), prod_imag.func(x))
\end{python}
\end{minipage}
\end{center}

Splitting the real and imaginary parts has the advantage of keeping both of 
them real and therefore also their integrals. This allows to make full usage of 
the \zfit{} integration. Since these parts are 
\textit{independent} of any coefficient, this allows to cache the 
integral value.

\begin{center}
\begin{minipage}{\textwidth}
\begin{python}
class AmplitudeProductProjection(zfit.func.BaseFunctorFunc):
	def __init__(self, amp1: ZfitFunc, amp2: ZfitFunc, proj, ...):
	self.projector = proj
	super().__init__(funcs=[amp1, amp2], ...)
	self._cache_integral = None
	
	def _func(self, x):
		amp1, amp2 = self.funcs
		return self.projector(amp1.func(x) * tf.conj(amp2.func(x)))
		
	def _single_hook_integrate(self, limits, norm_range, name):
		integral = self._cache_integral
		if integral is None:
			integral = super()._single_hook_integrate(limits=limits, 
													  norm_range=norm_range,
													  name=name)
			integral = zfit.run(integral)  # simplified
			self._cache_integral = integral
		return integral
\end{python}
\end{minipage}
\end{center}
where we used the hook for the integration as described in Appendix 
\ref{appendix:basemodel}.

All that is left now are the coefficients, which are just parameters, and the 
resonances. Using the
Breit-Wigner function from 
\zphys{},\footnote{The Breit-Wigner is currently an open merge request.} we can 
build them

\begin{center}
\begin{minipage}{\textwidth}
\begin{python}
resonances = [
	'rho(770)': RelativisticBreitWigner(rho_770_plus_mass,
										obs=zfit.Space('m2pipi', limits...))
	...
	]
coeffs = [
	zfit.ComplexParameter.from_polar('c_rho770', 1.0, 0.0),
	...
	]
\end{python}
\end{minipage}
\end{center}
and combining all of the above we can finally create the whole model

\begin{center}
\begin{minipage}{\textwidth}
\begin{python}
obs = zfit.Space(['m2kpim', 'm2kpi0', 'm2pipi', limits=...])
pdf = SumAmplitudeSquaredPDF(obs=obs, coeffs=coeffs, amps=resonances, ...)
\end{python}
\end{minipage}
\end{center}

This PDF is quite an accomplishment: while the implementation compared to all 
previous examples is more extensive, \text{none} of the above 
parts is \textit{really} complicated to implement and uses basic 
\zfit{} functionalities. Nonetheless, with this few lines of code, \zfit{} 
extends its functionality into the world of amplitude fitters, where usually 
only dedicated tools exist.

A missing part is the sampling: since the kinematics of the particles is not 
encoded into the PDF, it is needed inside the sampling. For that, we can create 
a decay as shown in examples above in Sec. \ref{py:reso_particle}  with the 
\zphasespace{} package. For simplicity, the Breit-Wigner distribution is 
used to model the mass shape. 
Building the decay and sample from it can be registered with the PDF as 
importance sampling.\footnote{Since this feature is not yet fully public in 
\zfit{}, no explicit example is shown. The mechanism will be similar to the 
registration of an integral.}

\begin{figure}[tbp]
\label{fig:dalitz_comparison}
\begin{subfigure}{.43\textwidth}
\includegraphics[width=\textwidth]{figs/dalitz_paper_cut.png}
% \caption{}
\end{subfigure}%
\begin{subfigure}{.57\textwidth}
\includegraphics[width=\textwidth]{figs/dalitz_kpim_kpiz.png}
% \caption{}
\end{subfigure}%
\caption{Dalitz plot of the invariant mass $\Kp\pim$ and $\Kp\piz$. The left 
histogram is taken from the paper while the one on the right was sampled from a 
distribution built with \zfit{} and using \zphasespace{} as sampler.}
\end{figure}



In order to compare the \zfit{} implementation with the measured data, a sample 
is drawn from the PDF. It is to note that the right plot in Fig. 
\ref{fig:dalitz_comparison}, drawn with 
\zfit{}, is currently not representative, since an instability in the sampling 
can vary the samples strongly depending on a single proposed 
event.\footnote{As described in \ref{sec:sampling techniques}, weights going 
	to zero as occurring in the phasespace can cause large problems.}. While 
	the 
two distributions 
do not agree, it is notable that resonances are 
taken into account and that the phasespace is sampled with correct borders. 
This unveils another current limitation in \zfit{}, that currently only 
rectangular limits are possible. As a future extension, arbitrary limits for 
\zspace{} will be implemented.

What was shown in this example are the low level components to build an 
amplitude analysis using \zfit{}. Most of the work above is straight forward 
but cumbersome 
and can be hidden behind a higher level interface that takes care of the right 
assignment of observables, resonances and compositions. Furthermore, and adding 
an additional 
layer around the resonances, the \zphasespace{} decays can be kept with the 
actual 
resonances which allows the phasespace to be an integral part of the amplitude. 
A higher level interface, which leaves the user with the 
specification of the resonances and coefficients but still offers the 
possibility to 
replace any part that was shown above, is currently under work in \zamp.