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