# Copyright (c) 2023 zfit
"""Recurrent polynomials."""
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
if TYPE_CHECKING:
import zfit
from collections.abc import Mapping
import abc
import tensorflow as tf
import zfit.z.numpy as znp
from zfit import z
from ..core.basepdf import BasePDF
from ..core.space import Space
from ..settings import ztypes
from ..util import ztyping
from ..util.container import convert_to_container
from ..util.exception import SpecificFunctionNotImplemented
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: tf.Tensor | None = None,
*,
extended: ztyping.ParamTypeInput | None = None,
name: str = "Polynomial",
): # noqa
"""Base class to create 1 dimensional recursive polynomials that can be rescaled. Overwrite _poly_func.
Args:
obs: |@doc:pdf.init.obs||@docend:pdf.init.obs|
coeffs: |@doc:pdf.polynomial.init.coeffs| Coefficients of the sum of the polynomial.
The coefficients of the polynomial, starting with the first order
term. To set the constant term, use ``coeff0``. |@docend:pdf.polynomial.init.coeffs|
apply_scaling: |@doc:pdf.polynomial.init.apply_scaling| Rescale the data so that the actual limits represent (-1, 1).
This is usually wanted as the polynomial is defined in this range.
Default is ``True``. |@docend:pdf.polynomial.init.apply_scaling|
.. math::
x_{n+1} = recurrence(x_{n}, x_{n-1}, n)
coeff0: |@doc:pdf.polynomial.init.coeff0| Coefficient of the constant term.
This is the coefficient of the constant term, i.e. the term
:math:`x^0`. If None, set to 1. |@docend:pdf.polynomial.init.coeff0|
extended: |@doc:pdf.init.extended| The overall yield of the PDF.
If this is parameter-like, it will be used as the yield,
the expected number of events, and the PDF will be extended.
An extended PDF has additional functionality, such as the
``ext_*`` methods and the ``counts`` (for binned PDFs). |@docend:pdf.init.extended|
name: |@doc:pdf.init.name||@docend:pdf.init.name|
"""
# 0th coefficient set to 1 by default
coeff0 = (
z.constant(1.0) 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, extended=extended)
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 SpecificFunctionNotImplemented
def create_poly(x, polys, coeffs, recurrence):
degree = len(coeffs) - 1
polys = do_recurrence(x, polys=polys, degree=degree, recurrence=recurrence)
sum_polys = znp.sum([coeff * poly for coeff, poly in zip(coeffs, polys)], axis=0)
return sum_polys
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", stateless_args=True)
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) * znp.multiply(x, p1) - n * p2) / (n + 1)
def legendre_shape(x, coeffs):
return create_poly(
x=x, polys=legendre_polys, coeffs=coeffs, recurrence=legendre_recurrence
)
def legendre_integral(
limits: ztyping.SpaceType,
norm: 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["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.0 * (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 = znp.reshape(integral, newshape=())
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: ztyping.ParamTypeInput | None = None,
*,
extended: ztyping.ParamTypeInput | None = 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: |@doc:pdf.init.obs||@docend:pdf.init.obs|
coeffs: |@doc:pdf.init.coeffs||@docend:pdf.init.coeffs|
apply_scaling: |@doc:pdf.init.apply_scaling||@docend:pdf.init.apply_scaling|
coeff0: |@doc:pdf.init.coeff0||@docend:pdf.init.coeff0|
name: |@doc:pdf.init.name||@docend:pdf.init.name|
extended: |@doc:pdf.init.extended| The overall yield of the PDF.
If this is parameter-like, it will be used as the yield,
the expected number of events, and the PDF will be extended.
An extended PDF has additional functionality, such as the
``ext_*`` methods and the ``counts`` (for binned PDFs). |@docend:pdf.init.extended|
"""
super().__init__(
obs=obs,
name=name,
coeffs=coeffs,
apply_scaling=apply_scaling,
coeff0=coeff0,
extended=extended,
)
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", stateless_args=True)
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 * znp.multiply(x, p1) - p2
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: ztyping.ParamTypeInput | None = None,
*,
extended: ztyping.ParamTypeInput | None = 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: |@doc:pdf.init.obs||@docend:pdf.init.obs|
coeffs: |@doc:pdf.init.coeffs||@docend:pdf.init.coeffs|
apply_scaling: |@doc:pdf.init.apply_scaling||@docend:pdf.init.apply_scaling|
coeff0: |@doc:pdf.init.coeff0||@docend:pdf.init.coeff0|
name: |@doc:pdf.init.name||@docend:pdf.init.name|
extended: |@doc:pdf.init.extended| The overall yield of the PDF.
If this is parameter-like, it will be used as the yield,
the expected number of events, and the PDF will be extended.
An extended PDF has additional functionality, such as the
``ext_*`` methods and the ``counts`` (for binned PDFs). |@docend:pdf.init.extended|
"""
super().__init__(
obs=obs,
name=name,
coeffs=coeffs,
coeff0=coeff0,
apply_scaling=apply_scaling,
extended=extended,
)
def _poly_func(self, x):
coeffs = convert_coeffs_dict_to_list(self.params)
return chebyshev_shape(x=x, coeffs=coeffs)
def func_integral_chebyshev1(limits, norm, 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["c_0"] * (
upper - lower
) # if polynomial 0 is defined as T_0 = 1
if model.degree >= 1:
integral += (
model.params["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 = znp.reshape(integral, newshape=())
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]
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: ztyping.ParamTypeInput | None = None,
*,
extended: ztyping.ParamTypeInput | None = 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: |@doc:pdf.init.obs||@docend:pdf.init.obs|
coeffs: |@doc:pdf.init.coeffs||@docend:pdf.init.coeffs|
apply_scaling: |@doc:pdf.init.apply_scaling||@docend:pdf.init.apply_scaling|
coeff0: |@doc:pdf.init.coeff0||@docend:pdf.init.coeff0|
name: |@doc:pdf.init.name||@docend:pdf.init.name|
extended: |@doc:pdf.init.extended| The overall yield of the PDF.
If this is parameter-like, it will be used as the yield,
the expected number of events, and the PDF will be extended.
An extended PDF has additional functionality, such as the
``ext_*`` methods and the ``counts`` (for binned PDFs). |@docend:pdf.init.extended|
"""
super().__init__(
obs=obs,
name=name,
coeffs=coeffs,
coeff0=coeff0,
apply_scaling=apply_scaling,
extended=extended,
)
def _poly_func(self, x):
coeffs = convert_coeffs_dict_to_list(self.params)
return chebyshev2_shape(x=x, coeffs=coeffs)
def func_integral_chebyshev2(limits, norm, 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.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 = znp.reshape(integral, newshape=())
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
)
def generalized_laguerre_polys_factory(alpha=0.0):
return [lambda x: tf.ones_like(x), lambda x: 1 + alpha - x]
laguerre_polys = generalized_laguerre_polys_factory(alpha=0.0)
def generalized_laguerre_recurrence_factory(alpha=0.0):
@z.function(wraps="zfit_tensor", stateless_args=True)
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 (znp.multiply(2 * n + 1 + alpha - x, p1) - (n + alpha) * p2) / (n + 1)
return generalized_laguerre_recurrence
laguerre_recurrence = generalized_laguerre_recurrence_factory(alpha=0.0)
def generalized_laguerre_shape_factory(alpha=0.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.0)
laguerre_shape_alpha_minusone = generalized_laguerre_shape_factory(
alpha=-1.0
) # for integral
[docs]class Laguerre(RecursivePolynomial):
def __init__(
self,
obs,
coeffs: list,
apply_scaling: bool = True,
coeff0: ztyping.ParamTypeInput | None = None,
*,
extended: ztyping.ParamTypeInput | None = 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: |@doc:pdf.init.obs||@docend:pdf.init.obs|
coeffs: |@doc:pdf.init.coeffs||@docend:pdf.init.coeffs|
apply_scaling: |@doc:pdf.init.apply_scaling||@docend:pdf.init.apply_scaling|
coeff0: |@doc:pdf.init.coeff0||@docend:pdf.init.coeff0|
name: |@doc:pdf.init.name||@docend:pdf.init.name|
extended: |@doc:pdf.init.extended| The overall yield of the PDF.
If this is parameter-like, it will be used as the yield,
the expected number of events, and the PDF will be extended.
An extended PDF has additional functionality, such as the
``ext_*`` methods and the ``counts`` (for binned PDFs). |@docend:pdf.init.extended|
"""
super().__init__(
obs=obs,
name=name,
coeffs=coeffs,
coeff0=coeff0,
apply_scaling=apply_scaling,
extended=extended,
)
def _poly_func(self, x):
coeffs = convert_coeffs_dict_to_list(self.params)
return laguerre_shape(x=x, coeffs=coeffs)
def func_integral_laguerre(limits, norm, 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:
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.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 = znp.reshape(integral, newshape=())
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", stateless_args=True)
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 * (znp.multiply(x, p1) - n * p2)
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: ztyping.ParamTypeInput | None = None,
*,
extended: ztyping.ParamTypeInput | None = 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: |@doc:pdf.init.obs||@docend:pdf.init.obs|
coeffs: |@doc:pdf.init.coeffs||@docend:pdf.init.coeffs|
apply_scaling: |@doc:pdf.init.apply_scaling||@docend:pdf.init.apply_scaling|
coeff0: |@doc:pdf.init.coeff0||@docend:pdf.init.coeff0|
name: |@doc:pdf.init.name||@docend:pdf.init.name|
extended: |@doc:pdf.init.extended| The overall yield of the PDF.
If this is parameter-like, it will be used as the yield,
the expected number of events, and the PDF will be extended.
An extended PDF has additional functionality, such as the
``ext_*`` methods and the ``counts`` (for binned PDFs). |@docend:pdf.init.extended|
"""
super().__init__(
obs=obs,
name=name,
coeffs=coeffs,
coeff0=coeff0,
apply_scaling=apply_scaling,
extended=extended,
)
def _poly_func(self, x):
coeffs = convert_coeffs_dict_to_list(self.params)
return hermite_shape(x=x, coeffs=coeffs)
def func_integral_hermite(limits, norm, 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.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.0, 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 = znp.reshape(integral, newshape=())
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
)
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