Source code for zfit.models.polynomials

#  Copyright (c) 2020 zfit
"""Recurrent polynomials."""
import abc
from typing import List, Dict, Optional, Mapping

import tensorflow as tf

from zfit import z
from ..util import ztyping
from ..util.container import convert_to_container
from ..settings import ztypes
from ..core.space import Space
from ..core.basepdf import BasePDF
from ..util.exception import SpecificFunctionNotImplementedError

[docs]def rescale_minus_plus_one(x: tf.Tensor, limits: "zfit.Space") -> tf.Tensor:
"""Rescale and shift *x* as *limits* were rescaled and shifted to be in (-1, 1). Useful for orthogonal polynomials.

Args:
x: Array like data
limits: 1-D limits

Returns:
The rescaled tensor.
"""
lim_low, lim_high = limits.limit1d
x = (2 * x - lim_low - lim_high) / (lim_high - lim_low)
return x

[docs]class RecursivePolynomial(BasePDF):
"""1D polynomial generated via three-term recurrence.

"""

def __init__(self, obs, coeffs: list,
apply_scaling: bool = True, coeff0: Optional[tf.Tensor] = None,
name: str = "Polynomial"):  # noqa
"""Base class to create 1 dimensional recursive polynomials that can be rescaled. Overwrite _poly_func.

Args:
coeffs: Coefficients for each polynomial. Used to calculate the degree.
apply_scaling: Rescale the data so that the actual limits represent (-1, 1).

.. math::
x_{n+1} = recurrence(x_{n}, x_{n-1}, n)

"""
# 0th coefficient set to 1 by default
coeff0 = z.constant(1.) if coeff0 is None else tf.cast(coeff0, dtype=ztypes.float)
coeffs = convert_to_container(coeffs).copy()
coeffs.insert(0, coeff0)
params = {f"c_{i}": coeff for i, coeff in enumerate(coeffs)}
self._degree = len(coeffs) - 1  # 1 coeff -> 0th degree
self._do_scale = apply_scaling
if apply_scaling and not (isinstance(obs, Space) and obs.n_limits == 1):
raise ValueError("obs need to be a Space with exactly one limit if rescaling is requested.")
super().__init__(obs=obs, name=name, params=params)

def _polynomials_rescale(self, x):
if self._do_scale:
x = rescale_minus_plus_one(x, limits=self.space)
return x

@property
def degree(self):
"""int: degree of the polynomial, starting from 0."""
return self._degree

def _unnormalized_pdf(self, x):
x = x.unstack_x()
x = self._polynomials_rescale(x)
return self._poly_func(x=x)

@abc.abstractmethod
def _poly_func(self, x):
raise SpecificFunctionNotImplementedError

[docs]def create_poly(x, polys, coeffs, recurrence):
degree = len(coeffs) - 1
polys = do_recurrence(x, polys=polys, degree=degree, recurrence=recurrence)
sum_polys = tf.reduce_sum(input_tensor=[coeff * poly for coeff, poly in zip(coeffs, polys)], axis=0)
return sum_polys

[docs]def do_recurrence(x, polys, degree, recurrence):
polys = [polys[0](x), polys[1](x)]
for i_deg in range(1, degree):  # recurrence returns for the n+1th degree
polys.append(recurrence(polys[-1], polys[-2], i_deg, x))
return polys

legendre_polys = [lambda x: tf.ones_like(x), lambda x: x]

@z.function(wraps='zfit_tensor')
def legendre_recurrence(p1, p2, n, x):
"""Recurrence relation for Legendre polynomials.

.. math::
(n+1) P_{n+1}(x) = (2n + 1) x P_{n}(x) - n P_{n-1}(x)

"""
return ((2 * n + 1) * tf.multiply(x, p1) - n * p2) / (n + 1)

[docs]def legendre_shape(x, coeffs):
return create_poly(x=x, polys=legendre_polys, coeffs=coeffs, recurrence=legendre_recurrence)

[docs]def legendre_integral(limits: ztyping.SpaceType, norm_range: ztyping.SpaceType,
params: List["zfit.Parameter"], model: RecursivePolynomial):
"""Recursive integral of Legendre polynomials"""
lower, upper = limits.limit1d
lower_rescaled = model._polynomials_rescale(lower)
upper_rescaled = model._polynomials_rescale(upper)
# if np.allclose((lower_rescaled, upper_rescaled), (-1, 1)):
#     return z.constant(2.)  #

lower = z.convert_to_tensor(lower_rescaled)
upper = z.convert_to_tensor(upper_rescaled)

integral_0 = model.params[f"c_0"] * (upper - lower)  # if polynomial 0 is 1
if model.degree == 0:
integral = integral_0
else:

def indefinite_integral(limits):
max_degree = model.degree + 1  # needed +1 for integral, max poly in term for n is n+1
polys = do_recurrence(x=limits, polys=legendre_polys, degree=max_degree,
recurrence=legendre_recurrence)
one_limit_integrals = []
for degree in range(1, max_degree):
coeff = model.params[f"c_{degree}"]
one_limit_integrals.append(coeff * (polys[degree + 1] - polys[degree - 1]) /
(2. * (z.convert_to_tensor(degree)) + 1))
return z.reduce_sum(one_limit_integrals, axis=0)

integral = indefinite_integral(upper) - indefinite_integral(lower) + integral_0
integral = tf.reshape(integral, shape=())
integral *= 0.5 * model.space.area()  # rescale back to whole width

return integral

[docs]class Legendre(RecursivePolynomial):

def __init__(self, obs: ztyping.ObsTypeInput, coeffs: List[ztyping.ParamTypeInput],
apply_scaling: bool = True, coeff0: Optional[ztyping.ParamTypeInput] = None,
name: str = "Legendre"):  # noqa
"""Linear combination of Legendre polynomials of order len(coeffs), the coeffs are overall scaling factors.

The 0th coefficient is set to 1 by default but can be explicitly set with *coeff0*. Since the PDF normalization
removes a degree of freedom, the 0th coefficient is redundant and leads to an arbitrary overall scaling of all
parameters.

Notice that this is already a sum of polynomials and the coeffs are simply scaling the individual orders of the
polynomials.

The recursive definition of **a single order** of the polynomial is

.. math::
(n+1) P_{n+1}(x) = (2n + 1) x P_{n}(x) - n P_{n-1}(x)

with
P_0 = 1
P_1 = x

Args:
obs: The default space the PDF is defined in.
coeffs: A list of the coefficients for the polynomials of order 1+ in the sum.
apply_scaling: Rescale the data so that the actual limits represent (-1, 1).
coeff0: The scaling factor of the 0th order polynomial. If not given, it is set to 1.
name: Name of the polynomial
"""
super().__init__(obs=obs, name=name,
coeffs=coeffs, apply_scaling=apply_scaling, coeff0=coeff0)

def _poly_func(self, x):
coeffs = convert_coeffs_dict_to_list(self.params)
return legendre_shape(x=x, coeffs=coeffs)

legendre_limits = Space(axes=0, limits=(Space.ANY_LOWER, Space.ANY_UPPER))
Legendre.register_analytic_integral(func=legendre_integral, limits=legendre_limits)

chebyshev_polys = [lambda x: tf.ones_like(x), lambda x: x]

@z.function(wraps='zfit_tensor')
def chebyshev_recurrence(p1, p2, _, x):
"""Recurrence relation for Chebyshev polynomials.

T_{n+1}(x) = 2 x T_{n}(x) - T_{n-1}(x)

"""
return 2 * tf.multiply(x, p1) - p2

[docs]def chebyshev_shape(x, coeffs):
return create_poly(x=x, polys=chebyshev_polys, coeffs=coeffs, recurrence=chebyshev_recurrence)

[docs]class Chebyshev(RecursivePolynomial):

def __init__(self, obs, coeffs: list, apply_scaling: bool = True, coeff0: Optional[ztyping.ParamTypeInput] = None,
name: str = "Chebyshev"):  # noqa
"""Linear combination of Chebyshev (first kind) polynomials of order len(coeffs), coeffs are scaling factors.

The 0th coefficient is set to 1 by default but can be explicitly set with *coeff0*. Since the PDF normalization
removes a degree of freedom, the 0th coefficient is redundant and leads to an arbitrary overall scaling of all
parameters.

Notice that this is already a sum of polynomials and the coeffs are simply scaling the individual orders of the
polynomials.

The recursive definition of **a single order** of the polynomial is

.. math::
T_{n+1}(x) = 2 x T_{n}(x) - T_{n-1}(x)

with
T_{0} = 1
T_{1} = x

Notice that :math:T_1 is x as opposed to 2x in Chebyshev polynomials of the second kind.

Args:
obs: The default space the PDF is defined in.
coeffs: A list of the coefficients for the polynomials of order 1+ in the sum.
apply_scaling: Rescale the data so that the actual limits represent (-1, 1).
coeff0: The scaling factor of the 0th order polynomial. If not given, it is set to 1.
name: Name of the polynomial
"""
super().__init__(obs=obs, name=name,
coeffs=coeffs, coeff0=coeff0,
apply_scaling=apply_scaling)

def _poly_func(self, x):
coeffs = convert_coeffs_dict_to_list(self.params)
return chebyshev_shape(x=x, coeffs=coeffs)

[docs]def func_integral_chebyshev1(limits, norm_range, params, model):
lower, upper = limits.rect_limits
lower_rescaled = model._polynomials_rescale(lower)
upper_rescaled = model._polynomials_rescale(upper)

lower = z.convert_to_tensor(lower_rescaled)
upper = z.convert_to_tensor(upper_rescaled)

integral = model.params[f"c_0"] * (upper - lower)  # if polynomial 0 is defined as T_0 = 1
if model.degree >= 1:
integral += model.params[f"c_1"] * 0.5 * (upper ** 2 - lower ** 2)  # if polynomial 0 is defined as T_0 = 1
if model.degree >= 2:

def indefinite_integral(limits):
max_degree = model.degree + 1
polys = do_recurrence(x=limits, polys=chebyshev_polys, degree=max_degree,
recurrence=chebyshev_recurrence)
one_limit_integrals = []
for degree in range(2, max_degree):
coeff = model.params[f"c_{degree}"]
n_float = z.convert_to_tensor(degree)
integral = (n_float * polys[degree + 1] / (z.square(n_float) - 1)
- limits * polys[degree] / (n_float - 1))
one_limit_integrals.append(coeff * integral)
return z.reduce_sum(one_limit_integrals, axis=0)

integral += indefinite_integral(upper) - indefinite_integral(lower)
integral = tf.reshape(integral, shape=())
integral *= 0.5 * model.space.area()  # rescale back to whole width
integral = tf.gather(integral, indices=0, axis=-1)
return integral

chebyshev1_limits_integral = Space(axes=0, limits=(Space.ANY_LOWER, Space.ANY_UPPER))
Chebyshev.register_analytic_integral(func=func_integral_chebyshev1, limits=chebyshev1_limits_integral)

chebyshev2_polys = [lambda x: tf.ones_like(x), lambda x: x * 2]

[docs]def chebyshev2_shape(x, coeffs):
return create_poly(x=x, polys=chebyshev2_polys, coeffs=coeffs, recurrence=chebyshev_recurrence)

[docs]class Chebyshev2(RecursivePolynomial):

def __init__(self, obs, coeffs: list, apply_scaling: bool = True, coeff0: Optional[ztyping.ParamTypeInput] = None,
name: str = "Chebyshev2"):  # noqa
"""Linear combination of Chebyshev (second kind) polynomials of order len(coeffs), coeffs are scaling factors.

The 0th coefficient is set to 1 by default but can be explicitly set with *coeff0*. Since the PDF normalization
removes a degree of freedom, the 0th coefficient is redundant and leads to an arbitrary overall scaling of all
parameters.

Notice that this is already a sum of polynomials and the coeffs are simply scaling the individual orders of the
polynomials.

The recursive definition of **a single order** of the polynomial is

.. math::
T_{n+1}(x) = 2 x T_{n}(x) - T_{n-1}(x)

with
T_{0} = 1
T_{1} = 2x

Notice that :math:T_1 is 2x as opposed to x in Chebyshev polynomials of the first kind.

Args:
obs: The default space the PDF is defined in.
coeffs: A list of the coefficients for the polynomials of order 1+ in the sum.
apply_scaling: Rescale the data so that the actual limits represent (-1, 1).
coeff0: The scaling factor of the 0th order polynomial. If not given, it is set to 1.
name: Name of the polynomial
"""
super().__init__(obs=obs, name=name,
coeffs=coeffs, coeff0=coeff0, apply_scaling=apply_scaling)

def _poly_func(self, x):
coeffs = convert_coeffs_dict_to_list(self.params)
return chebyshev2_shape(x=x, coeffs=coeffs)

[docs]def func_integral_chebyshev2(limits, norm_range, params, model):
lower, upper = limits.limit1d
lower_rescaled = model._polynomials_rescale(lower)
upper_rescaled = model._polynomials_rescale(upper)

lower = z.convert_to_tensor(lower_rescaled)
upper = z.convert_to_tensor(upper_rescaled)

# the integral of cheby2_ni is a cheby1_ni+1/(n+1). We add the (n+1) to the coeffs. The cheby1 shape makes
# the sum for us.
coeffs_cheby1 = {'c_0': z.constant(0., dtype=model.dtype)}

for name, coeff in params.items():
n_plus1 = int(name.split("_", 1)[-1]) + 1
coeffs_cheby1[f'c_{n_plus1}'] = coeff / z.convert_to_tensor(n_plus1, dtype=model.dtype)
coeffs_cheby1 = convert_coeffs_dict_to_list(coeffs_cheby1)

def indefinite_integral(limits):
return chebyshev_shape(x=limits, coeffs=coeffs_cheby1)

integral = indefinite_integral(upper) - indefinite_integral(lower)
integral = tf.reshape(integral, shape=())
integral *= 0.5 * model.space.area()  # rescale back to whole width

return integral

chebyshev2_limits_integral = Space(axes=0, limits=(Space.ANY_LOWER, Space.ANY_UPPER))
Chebyshev2.register_analytic_integral(func=func_integral_chebyshev2, limits=chebyshev2_limits_integral)

[docs]def generalized_laguerre_polys_factory(alpha=0.):
return [lambda x: tf.ones_like(x), lambda x: 1 + alpha - x]

laguerre_polys = generalized_laguerre_polys_factory(alpha=0.)

[docs]def generalized_laguerre_recurrence_factory(alpha=0.):
@z.function(wraps='zfit_tensor')
def generalized_laguerre_recurrence(p1, p2, n, x):
"""Recurrence relation for Laguerre polynomials.

:math:(n+1) L_{n+1}(x) = (2n + 1 + \alpha - x) L_{n}(x) - (n + \alpha) L_{n-1}(x)

"""
return (tf.multiply(2 * n + 1 + alpha - x, p1) - (n + alpha) * p2) / (n + 1)

return generalized_laguerre_recurrence

laguerre_recurrence = generalized_laguerre_recurrence_factory(alpha=0.)

[docs]def generalized_laguerre_shape_factory(alpha=0.):
recurrence = generalized_laguerre_recurrence_factory(alpha=alpha)
polys = generalized_laguerre_polys_factory(alpha=alpha)

def general_laguerre_shape(x, coeffs):
return create_poly(x=x, polys=polys, coeffs=coeffs, recurrence=recurrence)

return general_laguerre_shape

laguerre_shape = generalized_laguerre_shape_factory(alpha=0.)
laguerre_shape_alpha_minusone = generalized_laguerre_shape_factory(alpha=-1.)  # for integral

[docs]class Laguerre(RecursivePolynomial):

def __init__(self, obs, coeffs: list, apply_scaling: bool = True, coeff0: Optional[ztyping.ParamTypeInput] = None,
name: str = "Laguerre"):  # noqa
"""Linear combination of Laguerre polynomials of order len(coeffs), the coeffs are overall scaling factors.

The 0th coefficient is set to 1 by default but can be explicitly set with *coeff0*. Since the PDF normalization
removes a degree of freedom, the 0th coefficient is redundant and leads to an arbitrary overall scaling of all
parameters.

Notice that this is already a sum of polynomials and the coeffs are simply scaling the individual orders of the
polynomials.

The recursive definition of **a single order** of the polynomial is

.. math::

(n+1) L_{n+1}(x) = (2n + 1 + \alpha - x) L_{n}(x) - (n + \alpha) L_{n-1}(x)

with
P_0 = 1
P_1 = 1 - x

Args:
obs: The default space the PDF is defined in.
coeffs: A list of the coefficients for the polynomials of order 1+ in the sum.
apply_scaling: Rescale the data so that the actual limits represent (-1, 1).
coeff0: The scaling factor of the 0th order polynomial. If not given, it is set to 1.
name: Name of the polynomial
"""
super().__init__(obs=obs, name=name, coeffs=coeffs, coeff0=coeff0,
apply_scaling=apply_scaling)

def _poly_func(self, x):
coeffs = convert_coeffs_dict_to_list(self.params)
return laguerre_shape(x=x, coeffs=coeffs)

[docs]def func_integral_laguerre(limits, norm_range, params: Dict, model):
"""The integral of the simple laguerre polynomials.

Defined as :math:\int L_{n} = (-1) L_{n+1}^{(-1)} with :math:L^{(\alpha)} the generalized Laguerre polynom.

Args:
limits:
norm_range:
params:
model:

Returns:

"""
lower, upper = limits.limit1d
lower_rescaled = model._polynomials_rescale(lower)
upper_rescaled = model._polynomials_rescale(upper)

lower = z.convert_to_tensor(lower_rescaled)
upper = z.convert_to_tensor(upper_rescaled)

# The laguerre shape makes the sum for us. setting the 0th coeff to 0, since no -1 term exists.
coeffs_laguerre_nup = {f'c_{int(n.split("_", 1)[-1]) + 1}': c
for i, (n, c) in enumerate(params.items())}  # increase n -> n+1 of naming
coeffs_laguerre_nup['c_0'] = tf.constant(0., dtype=model.dtype)
coeffs_laguerre_nup = convert_coeffs_dict_to_list(coeffs_laguerre_nup)

def indefinite_integral(limits):
return -1 * laguerre_shape_alpha_minusone(x=limits, coeffs=coeffs_laguerre_nup)

integral = indefinite_integral(upper) - indefinite_integral(lower)
integral = tf.reshape(integral, shape=())
integral *= 0.5 * model.space.area()  # rescale back to whole width
return integral

laguerre_limits_integral = Space(axes=0, limits=(Space.ANY_LOWER, Space.ANY_UPPER))
Laguerre.register_analytic_integral(func=func_integral_laguerre, limits=laguerre_limits_integral)

hermite_polys = [lambda x: tf.ones_like(x), lambda x: 2 * x]

@z.function(wraps='zfit_tensor')
def hermite_recurrence(p1, p2, n, x):
"""Recurrence relation for Hermite polynomials (physics).

:math:H_{n+1}(x) = 2x H_{n}(x) - 2n H_{n-1}(x)

"""
return 2 * (tf.multiply(x, p1) - n * p2)

[docs]def hermite_shape(x, coeffs):
return create_poly(x=x, polys=hermite_polys, coeffs=coeffs, recurrence=hermite_recurrence)

[docs]class Hermite(RecursivePolynomial):

def __init__(self, obs, coeffs: list, apply_scaling: bool = True, coeff0: Optional[ztyping.ParamTypeInput] = None,
name: str = "Hermite"):  # noqa
"""Linear combination of Hermite polynomials (for physics) of order len(coeffs), with coeffs as scaling factors.

The 0th coefficient is set to 1 by default but can be explicitly set with *coeff0*. Since the PDF normalization
removes a degree of freedom, the 0th coefficient is redundant and leads to an arbitrary overall scaling of all
parameters.

Notice that this is already a sum of polynomials and the coeffs are simply scaling the individual orders of the
polynomials.

The recursive definition of **a single order** of the polynomial is

.. math::

H_{n+1}(x) = 2x H_{n}(x) - 2n H_{n-1}(x)

with
P_0 = 1
P_1 = 2x

Args:
obs: The default space the PDF is defined in.
coeffs: A list of the coefficients for the polynomials of order 1+ in the sum.
apply_scaling: Rescale the data so that the actual limits represent (-1, 1).
coeff0: The scaling factor of the 0th order polynomial. If not given, it is set to 1.
name: Name of the polynomial
"""
super().__init__(obs=obs, name=name, coeffs=coeffs, coeff0=coeff0,
apply_scaling=apply_scaling)

def _poly_func(self, x):
coeffs = convert_coeffs_dict_to_list(self.params)
return hermite_shape(x=x, coeffs=coeffs)

[docs]def func_integral_hermite(limits, norm_range, params, model):
lower, upper = limits.limit1d
lower_rescaled = model._polynomials_rescale(lower)
upper_rescaled = model._polynomials_rescale(upper)

lower = z.convert_to_tensor(lower_rescaled)
upper = z.convert_to_tensor(upper_rescaled)

# the integral of hermite is a hermite_ni. We add the ni to the coeffs.
coeffs = {'c_0': z.constant(0., dtype=model.dtype)}

for name, coeff in params.items():
ip1_coeff = int(name.split("_", 1)[-1]) + 1
coeffs[f'c_{ip1_coeff}'] = coeff / z.convert_to_tensor(ip1_coeff * 2., dtype=model.dtype)
coeffs = convert_coeffs_dict_to_list(coeffs)

def indefinite_integral(limits):
return hermite_shape(x=limits, coeffs=coeffs)

integral = indefinite_integral(upper) - indefinite_integral(lower)
integral = tf.reshape(integral, shape=())
integral *= 0.5 * model.space.area()  # rescale back to whole width

return integral

hermite_limits_integral = Space(axes=0, limits=(Space.ANY_LOWER, Space.ANY_UPPER))
Hermite.register_analytic_integral(func=func_integral_hermite, limits=hermite_limits_integral)

[docs]def convert_coeffs_dict_to_list(coeffs: Mapping) -> List:
# HACK(Mayou36): how to solve elegantly? yield not a param, only a dependent?
coeffs_list = []
for i in range(len(coeffs)):
try:
coeffs_list.append(coeffs[f"c_{i}"])
except KeyError:  # happens, if there are other parameters in there, such as a yield
break
return coeffs_list

# EOF