# Source code for zfit._loss.binnedloss

#  Copyright (c) 2024 zfit

from __future__ import annotations

from collections.abc import Iterable

import numpy as np
import tensorflow as tf
from uhi.typing.plottable import PlottableHistogram

from .. import z
from ..core.interfaces import ZfitBinnedData, ZfitBinnedPDF, ZfitParameter
from ..core.loss import BaseLoss
from ..util import ztyping
from ..util.checks import NONE
from ..util.container import convert_to_container
from ..util.warnings import warn_advanced_feature
from ..util.ztyping import OptionsInputType, ConstraintsInputType
from ..z import numpy as znp

@z.function(wraps="tensor", keepalive=True)
def _spd_transform(values, probs, variances):
"""Transform the data to the SPD form.

Scaled Poisson distribution from Bohm and Zech, NIMA 748 (2014) 1-6

The scale is >= 1 to ensure that empty data bins are handled correctly.

Args:
values: Data values, the counts in each bin.
probs: Data probabilities, the expected counts in each bin.
variances: Data variances, the variances of the counts in each bin.

Returns:
The transformed probabilities and values.
"""
# Scaled Poisson distribution from Bohm and Zech, NIMA 748 (2014) 1-6
scale = znp.maximum(
values * tf.math.reciprocal_no_nan(variances), znp.ones_like(values)
)
probs = probs * scale
values = values * scale
return probs, values

@z.function(wraps="tensor", keepalive=True)
def poisson_loss_calc(probs, values, log_offset=None, variances=None):
"""Calculate the Poisson log probability for a given set of data.

Args:
probs: Probabilities of the data, i.e., the expected number of events in each bin.
values: Values of the data, i.e., the number of events in each bin.
log_offset: Optional offset to be added to the loss. Useful for adding a constant to the loss to improve
numerical stability.
variances: (currently ignored)
Returns:
The Poisson log probability for the given data.
"""
# Optional variances of the data. If not None, the Poisson loss is calculated using the
#             scaled Poisson distribution from Bohm and Zech, NIMA 748 (2014) 1-6
if log_offset is None:
log_offset = False
use_offset = log_offset is not False
if (
False and variances is not None
):  # TODO: this gives very different uncertainties?  if fixed, rechange the docs
values, probs = _spd_transform(values, probs, variances=variances)
values += znp.asarray(1e-307, dtype=znp.float64)
probs += znp.asarray(1e-307, dtype=znp.float64)
poisson_term = tf.nn.log_poisson_loss(
values, znp.log(probs), compute_full_loss=not use_offset  # TODO: correct offset
)  # TODO: optimization?

# cross-check
# import tensorflow_probability as tfp
# poisson_dist = tfp.distributions.Poisson(rate=probs)
# poisson_term = -poisson_dist.log_prob(values)
if use_offset:
log_offset = znp.asarray(log_offset, dtype=znp.float64)
poisson_term += log_offset
return poisson_term

class BaseBinned(BaseLoss):
def __init__(
self,
model: ztyping.BinnedPDFInputType,
data: ztyping.BinnedDataInputType,
constraints: ConstraintsInputType = None,
options: OptionsInputType = None,
):
model = convert_to_container(model)
data = convert_to_container(data)
from zfit._data.binneddatav1 import BinnedData

data = [
(
BinnedData.from_hist(d)
if (
isinstance(d, PlottableHistogram)
and not isinstance(d, ZfitBinnedData)
)
else d
)
for d in data
]
not_binned_pdf = [mod for mod in model if not isinstance(mod, ZfitBinnedPDF)]
not_binned_data = [dat for dat in data if not isinstance(dat, ZfitBinnedData)]
not_binned_pdf_msg = (
"The following PDFs are not binned but need to be. They can be wrapped in an "
f"BinnedFromUnbinnedPDF. {not_binned_pdf} "
)
not_binned_data_msg = (
"The following datasets are not binned but need to be. They can be converted to a binned "
f"using the to_binned method. {not_binned_data}"
)
error_msg = ""
if not_binned_pdf:
error_msg += not_binned_pdf_msg
if not_binned_data:
error_msg += not_binned_data_msg
if error_msg:
raise ValueError(error_msg)

super().__init__(
model=model,
data=data,
constraints=constraints,
fit_range=None,
options=options,
)

def create_new(
self,
model: ztyping.BinnedPDFInputType = NONE,
data: ztyping.BinnedDataInputType = NONE,
constraints: ConstraintsInputType = NONE,
options: OptionsInputType = NONE,
):
r"""Create a new binned loss of this type. This is preferrable over creating a new instance in most cases.

Internals, such as certain optimizations will be shared and therefore the loss is made comparable.

If something is not given, it will be taken from the current loss.

Args:
model: |@doc:loss.binned.init.model| Binned PDF(s) that return the normalized probability
(rel_counts or counts) for
*data* under the given parameters.
If multiple model and data are given, they will be used
in the same order to do a simultaneous fit. |@docend:loss.binned.init.model|
data: |@doc:loss.binned.init.data| Binned dataset that will be given to the *model*.
If multiple model and data are given, they will be used
in the same order to do a simultaneous fit. |@docend:loss.binned.init.data|
constraints: |@doc:loss.init.constraints| Auxiliary measurements ("constraints")
that add a likelihood term to the loss.

.. math::
\mathcal{L}(\theta) = \mathcal{L}_{unconstrained} \prod_{i} f_{constr_i}(\theta)

Usually, an auxiliary measurement -- by its very nature -S  should only be added once
to the loss. zfit does not automatically deduplicate constraints if they are given
multiple times, leaving the freedom for arbitrary constructs.

Constraints can also be used to restrict the loss by adding any kinds of penalties. |@docend:loss.init.constraints|
options: |@doc:loss.init.options| Additional options (as a dict) for the loss.
Current possibilities include:

- 'subtr_const' (default True): subtract from each points
log probability density a constant that
is approximately equal to the average log probability
density in the very first evaluation before
the summation. This brings the initial loss value closer to 0 and increases,
especially for large datasets, the numerical stability.

The value will be stored ith 'subtr_const_value' and can also be given
directly.

The subtraction should not affect the minimum as the absolute
value of the NLL is meaningless. However,
with this switch on, one cannot directly compare
different likelihoods absolute value as the constant
may differ! Use create_new in order to have a comparable likelihood
between different losses or use the full argument in the value function
to calculate the full loss with all constants.

These settings may extend over time. In order to make sure that a loss is the
same under the same data, make sure to use create_new instead of instantiating
a new loss as the former will automatically overtake any relevant constants
and behavior. |@docend:loss.init.options|

Returns:
"""
if model is NONE:
model = self.model
if data is NONE:
data = self.data
if constraints is NONE:
constraints = self.constraints
if constraints is not None:
constraints = constraints.copy()
if options is NONE:
options = self._options
if isinstance(options, dict):
options = options.copy()
return type(self)(
model=model, data=data, constraints=constraints, options=options
)

[docs]
class ExtendedBinnedNLL(BaseBinned):
def __init__(
self,
model: ztyping.BinnedPDFInputType,
data: ztyping.BinnedDataInputType,
constraints: ConstraintsInputType = None,
options: OptionsInputType = None,
):
r"""Extended binned likelihood using the expected number of events per bin with a poisson probability.

The binned likelihood is defined as

.. math::
\mathcal{L} = \product \mathcal{poiss}(N_{modelbin_i}, N_{databin_i})
= N_{databin_i}^{N_{modelbin_i}} \frac{e^{- N_{databin_i}}}{N_{modelbin_i}!}

where :math:databin_i is the :math:i^{th} bin in the data and
:math:modelbin_i is the :math:i^{th} bin of the model, the expected counts.

|@doc:loss.init.explain.simultaneous| A simultaneous fit can be performed by giving one or more model, data, to the loss. The
length of each has to match the length of the others

.. math::
\mathcal{L}_{simultaneous}(\theta | {data_0, data_1, ..., data_n})
= \prod_{i} \mathcal{L}(\theta_i, data_i)

where :math:\theta_i is a set of parameters and
a subset of :math:\theta |@docend:loss.init.explain.simultaneous|

|@doc:loss.init.explain.negativelog| For optimization purposes, it is often easier
to minimize a function and to use a log transformation. The actual loss is given by

.. math::
\mathcal{L} = - \sum_{i}^{n} ln(f(\theta|x_i))

and therefore being called "negative log ..." |@docend:loss.init.explain.negativelog|

Args:
model: |@doc:loss.binned.init.model| Binned PDF(s) that return the normalized probability
(rel_counts or counts) for
*data* under the given parameters.
If multiple model and data are given, they will be used
in the same order to do a simultaneous fit. |@docend:loss.binned.init.model|
data: |@doc:loss.binned.init.data| Binned dataset that will be given to the *model*.
If multiple model and data are given, they will be used
in the same order to do a simultaneous fit. |@docend:loss.binned.init.data|
constraints: |@doc:loss.init.constraints| Auxiliary measurements ("constraints")
that add a likelihood term to the loss.

.. math::
\mathcal{L}(\theta) = \mathcal{L}_{unconstrained} \prod_{i} f_{constr_i}(\theta)

Usually, an auxiliary measurement -- by its very nature -S  should only be added once
to the loss. zfit does not automatically deduplicate constraints if they are given
multiple times, leaving the freedom for arbitrary constructs.

Constraints can also be used to restrict the loss by adding any kinds of penalties. |@docend:loss.init.constraints|
options: |@doc:loss.init.options| Additional options (as a dict) for the loss.
Current possibilities include:

- 'subtr_const' (default True): subtract from each points
log probability density a constant that
is approximately equal to the average log probability
density in the very first evaluation before
the summation. This brings the initial loss value closer to 0 and increases,
especially for large datasets, the numerical stability.

The value will be stored ith 'subtr_const_value' and can also be given
directly.

The subtraction should not affect the minimum as the absolute
value of the NLL is meaningless. However,
with this switch on, one cannot directly compare
different likelihoods absolute value as the constant
may differ! Use create_new in order to have a comparable likelihood
between different losses or use the full argument in the value function
to calculate the full loss with all constants.

These settings may extend over time. In order to make sure that a loss is the
same under the same data, make sure to use create_new instead of instantiating
a new loss as the former will automatically overtake any relevant constants
and behavior. |@docend:loss.init.options|
"""

# readd below if fixed
#     |@doc:loss.init.explain.spdtransform| A scaled Poisson
self._errordef = 0.5
super().__init__(
model=model, data=data, constraints=constraints, options=options
)

@z.function(wraps="loss", keepalive=True)
def _loss_func(
self,
model: Iterable[ZfitBinnedPDF],
data: Iterable[ZfitBinnedData],
fit_range,
constraints,
log_offset,
):
poisson_terms = []
for mod, dat in zip(model, data):
values = dat.values(  # TODO: right order of model and data?
# obs=mod.obs
)
variances = dat.variances()
probs = mod.counts(dat)
poisson_term = poisson_loss_calc(probs, values, log_offset, variances)
poisson_term_summed = znp.sum(poisson_term)
poisson_terms.append(poisson_term_summed)  # TODO: change None
nll = znp.sum(poisson_terms)

if constraints:
if (
log_offset is False
):  # we need to check identity, cannot do runtime conditional if jitted
log_offset_val = 0.0
else:
log_offset_val = log_offset
log_offset_val = znp.asarray(log_offset_val, dtype=znp.float64)
constraints = z.reduce_sum(
[c.value() - log_offset_val * len(c.get_params()) for c in constraints]
)
nll += constraints

return nll

@property
def is_extended(self):
return True

def _get_params(
self,
floating: bool | None = True,
is_yield: bool | None = None,
extract_independent: bool | None = True,
) -> set[ZfitParameter]:
return super()._get_params(floating, is_yield, extract_independent)

[docs]
class BinnedNLL(BaseBinned):
def __init__(
self,
model: ztyping.BinnedPDFInputType,
data: ztyping.BinnedDataInputType,
constraints: ConstraintsInputType = None,
options: OptionsInputType = None,
):
r"""Binned negative log likelihood.

The binned likelihood is the binned version of :py:class:~zfit.loss.UnbinnedNLL. It is defined as

.. math::
\\mathcal{L} = \\product \\mathcal{poiss}(N_{modelbin_i}, N_{databin_i}) = N_{databin_i}^{N_{modelbin_i}} \frac{e^{- N_{databin_i}}}{N_{modelbin_i}!}

where :math:databin_i is the :math:i^{th} bin in the data and
:math:modelbin_i is the :math:i^{th} bin of the model multiplied by the total number of events in data.

|@doc:loss.init.explain.simultaneous| A simultaneous fit can be performed by giving one or more model, data, to the loss. The
length of each has to match the length of the others

.. math::
\mathcal{L}_{simultaneous}(\theta | {data_0, data_1, ..., data_n})
= \prod_{i} \mathcal{L}(\theta_i, data_i)

where :math:\theta_i is a set of parameters and
a subset of :math:\theta |@docend:loss.init.explain.simultaneous|

|@doc:loss.init.explain.negativelog| For optimization purposes, it is often easier
to minimize a function and to use a log transformation. The actual loss is given by

.. math::
\mathcal{L} = - \sum_{i}^{n} ln(f(\theta|x_i))

and therefore being called "negative log ..." |@docend:loss.init.explain.negativelog|

Args:
model: |@doc:loss.binned.init.model| Binned PDF(s) that return the normalized probability
(rel_counts or counts) for
*data* under the given parameters.
If multiple model and data are given, they will be used
in the same order to do a simultaneous fit. |@docend:loss.binned.init.model|
data: |@doc:loss.binned.init.data| Binned dataset that will be given to the *model*.
If multiple model and data are given, they will be used
in the same order to do a simultaneous fit. |@docend:loss.binned.init.data|
constraints: |@doc:loss.init.constraints| Auxiliary measurements ("constraints")
that add a likelihood term to the loss.

.. math::
\mathcal{L}(\theta) = \mathcal{L}_{unconstrained} \prod_{i} f_{constr_i}(\theta)

Usually, an auxiliary measurement -- by its very nature -S  should only be added once
to the loss. zfit does not automatically deduplicate constraints if they are given
multiple times, leaving the freedom for arbitrary constructs.

Constraints can also be used to restrict the loss by adding any kinds of penalties. |@docend:loss.init.constraints|
options: |@doc:loss.init.options| Additional options (as a dict) for the loss.
Current possibilities include:

- 'subtr_const' (default True): subtract from each points
log probability density a constant that
is approximately equal to the average log probability
density in the very first evaluation before
the summation. This brings the initial loss value closer to 0 and increases,
especially for large datasets, the numerical stability.

The value will be stored ith 'subtr_const_value' and can also be given
directly.

The subtraction should not affect the minimum as the absolute
value of the NLL is meaningless. However,
with this switch on, one cannot directly compare
different likelihoods absolute value as the constant
may differ! Use create_new in order to have a comparable likelihood
between different losses or use the full argument in the value function
to calculate the full loss with all constants.

These settings may extend over time. In order to make sure that a loss is the
same under the same data, make sure to use create_new instead of instantiating
a new loss as the former will automatically overtake any relevant constants
and behavior. |@docend:loss.init.options|
"""

# readd below if fixed
#            |@doc:loss.init.explain.spdtransform| A scaled Poisson distribution is...
self._errordef = 0.5
super().__init__(
model=model, data=data, constraints=constraints, options=options
)
extended_pdfs = [pdf for pdf in self.model if pdf.is_extended]
if extended_pdfs and type(self) == BinnedNLL:
warn_advanced_feature(
f"Extended PDFs ({extended_pdfs}) are given to a normal BinnedNLL. "
f" This won't take the yield "
"into account and simply treat the PDFs as non-extended PDFs. To create an "
"extended NLL, use the ExtendedBinnedNLL.",
identifier="extended_in_BinnedNLL",
)

@z.function(wraps="loss", keepalive=True)
def _loss_func(
self,
model: Iterable[ZfitBinnedPDF],
data: Iterable[ZfitBinnedData],
fit_range,
constraints,
log_offset,
):
poisson_terms = []
for mod, dat in zip(model, data):
values = dat.values(  # TODO: right order of model and data?
# obs=mod.obs
)
variances = dat.variances()
probs = mod.rel_counts(dat)
probs *= znp.sum(values)
poisson_term = poisson_loss_calc(probs, values, log_offset, variances)
poisson_term_summed = znp.sum(poisson_term)
poisson_terms.append(poisson_term_summed)
nll = znp.sum(poisson_terms)

if constraints:
if (
log_offset is False
):  # we need to check identity, cannot do runtime conditional if jitted
log_offset_val = 0.0
else:
log_offset_val = log_offset
log_offset_val = znp.asarray(log_offset_val, dtype=znp.float64)
constraints = z.reduce_sum(
[c.value() - log_offset_val * len(c.get_params()) for c in constraints]
)
nll += constraints

return nll

@property
def is_extended(self):
return False

def _get_params(
self,
floating: bool | None = True,
is_yield: bool | None = None,
extract_independent: bool | None = True,
) -> set[ZfitParameter]:
if not self.is_extended:
is_yield = False  # the loss does not depend on the yields
return super()._get_params(floating, is_yield, extract_independent)

@z.function(wraps="tensor", keepalive=True)
def chi2_loss_calc(probs, values, variances, log_offset=None, ignore_empty=None):
"""Calculate the chi2 for a given set of data.

Args:
probs: Probabilities of the data, i.e. the expected number of events in each bin.
values: Values of the data, i.e. the number of events in each
bin.
variances: Data variances, the variances of the counts in each bin.
log_offset: Optional offset to be added to the loss. Useful for adding a constant to the loss to improve
numerical stability.
ignore_empty: If True, empty bins are ignored.
Returns:
The chi2 for the given data.
"""

if log_offset is None:
log_offset = False
if ignore_empty is None:
ignore_empty = True
chi2_term = tf.math.squared_difference(probs, values)
if ignore_empty:
one_over_var = tf.math.reciprocal_no_nan(variances)
else:
one_over_var = tf.math.reciprocal(variances)
chi2_term *= one_over_var
if log_offset is not False:
log_offset = znp.asarray(log_offset, dtype=znp.float64)
chi2_term += log_offset
chi2_term = znp.sum(chi2_term)
return chi2_term

def _check_small_counts_chi2(data, ignore_empty):
for dat in data:
variances = dat.variances()
smaller_than_six = dat.values() < 6
if variances is None:
raise ValueError(f"variances cannot be None for Chi2: {dat}")
elif np.any(variances <= 0) and not ignore_empty:
raise ValueError(
f"Variances of {dat} contains zeros or negative numbers, cannot calculate chi2."
f" {variances}"
)
elif np.any(smaller_than_six):
warn_advanced_feature(
f"Some values in {dat} are < 6, the chi2 assumption of gaussian distributed"
f" uncertainties most likely won't hold anymore. Use Chi2 for large samples."
f"For smaller samples, consider using (Extended)BinnedNLL (or an unbinned fit).",
identifier="chi2_counts_small",
)

[docs]
class BinnedChi2(BaseBinned):
def __init__(
self,
model: ztyping.BinnedPDFInputType,
data: ztyping.BinnedDataInputType,
constraints: ConstraintsInputType = None,
options: OptionsInputType = None,
):
r"""Binned Chi2 loss, using the :math:N_{tot} from the data.

.. math::
\chi^2 = \sum_{\mathrm{bins}} \left( \frac{N_\mathrm{PDF,bin} - N_\mathrm{Data,bin}}{\sigma_\mathrm{Data,bin}} \right)^2

where

.. math::
N_\mathrm{PDF,bin} = \mathrm{pdf}(\text{integral}) \cdot N_\mathrm{Data,tot}
\sigma_\mathrm{bin} = \text{variance}

with variance the value of :class:~zfit.data.BinnedData.variances of the binned data.

|@doc:loss.init.binned.explain.chi2zeros| If the dataset has empty bins, the errors
will be zero and :math:\chi^2 is undefined. Two possibilities are available and
can be given as an option:

- "empty": "ignore" will ignore all bins with zero entries and won't count to the loss
- "errors": "expected" will use the expected counts from the model
with a Poissonian uncertainty |@docend:loss.init.binned.explain.chi2zeros|

Args:
model: |@doc:loss.binned.init.model| Binned PDF(s) that return the normalized probability
(rel_counts or counts) for
*data* under the given parameters.
If multiple model and data are given, they will be used
in the same order to do a simultaneous fit. |@docend:loss.binned.init.model|
data: |@doc:loss.binned.init.data| Binned dataset that will be given to the *model*.
If multiple model and data are given, they will be used
in the same order to do a simultaneous fit. |@docend:loss.binned.init.data|
constraints: |@doc:loss.init.constraints| Auxiliary measurements ("constraints")
that add a likelihood term to the loss.

.. math::
\mathcal{L}(\theta) = \mathcal{L}_{unconstrained} \prod_{i} f_{constr_i}(\theta)

Usually, an auxiliary measurement -- by its very nature -S  should only be added once
to the loss. zfit does not automatically deduplicate constraints if they are given
multiple times, leaving the freedom for arbitrary constructs.

Constraints can also be used to restrict the loss by adding any kinds of penalties. |@docend:loss.init.constraints|
options: |@doc:loss.init.options| Additional options (as a dict) for the loss.
Current possibilities include:

- 'subtr_const' (default True): subtract from each points
log probability density a constant that
is approximately equal to the average log probability
density in the very first evaluation before
the summation. This brings the initial loss value closer to 0 and increases,
especially for large datasets, the numerical stability.

The value will be stored ith 'subtr_const_value' and can also be given
directly.

The subtraction should not affect the minimum as the absolute
value of the NLL is meaningless. However,
with this switch on, one cannot directly compare
different likelihoods absolute value as the constant
may differ! Use create_new in order to have a comparable likelihood
between different losses or use the full argument in the value function
to calculate the full loss with all constants.

These settings may extend over time. In order to make sure that a loss is the
same under the same data, make sure to use create_new instead of instantiating
a new loss as the former will automatically overtake any relevant constants
and behavior. |@docend:loss.init.options|
"""
self._errordef = 1.0
if options is None:
options = {}
if options.get("empty") is None:
options["empty"] = "ignore"
if options.get("errors") is None:
options["errors"] = "data"
super().__init__(
model=model, data=data, constraints=constraints, options=options
)
extended_pdfs = [pdf for pdf in self.model if pdf.is_extended]
if extended_pdfs and type(self) == BinnedChi2:
warn_advanced_feature(
f"Extended PDFs ({extended_pdfs}) are given to a normal BinnedChi2. "
f" This won't take the yield "
"into account and simply treat the PDFs as non-extended PDFs. To create an "
"extended loss, use the ExtendedBinnedChi2.",
identifier="extended_in_BinnedChi2",
)

def _precompile(self):
super()._precompile()
ignore_empty = (
self._options.get("empty") == "ignore"
or self._options.get("errors") == "expected"
)
data = self.data
_check_small_counts_chi2(data, ignore_empty)

@z.function(wraps="loss", keepalive=True)
def _loss_func(
self,
model: Iterable[ZfitBinnedPDF],
data: Iterable[ZfitBinnedData],
fit_range,
constraints,
log_offset,
):
del fit_range
ignore_empty = self._options.get("empty") == "ignore"
chi2_terms = []
if log_offset is False:
log_offset_val = 0.0
else:
log_offset_val = log_offset
log_offset_val = znp.asarray(log_offset_val, dtype=znp.float64)

for mod, dat in zip(model, data):
values = dat.values(  # TODO: right order of model and data?
# obs=mod.obs
)
probs = mod.rel_counts(dat)
probs *= znp.sum(values)

variance_method = self._options.get("errors")
if variance_method == "expected":
variances = znp.sqrt(probs + znp.asarray(1e-307, dtype=znp.float64))
elif variance_method == "data":
variances = dat.variances()
else:
raise ValueError()
if variances is None:
raise ValueError(f"variances cannot be None for Chi2: {dat}")

chi2_term = chi2_loss_calc(
probs, values, variances, log_offset, ignore_empty=ignore_empty
)
chi2_terms.append(chi2_term)
chi2_term = znp.sum(chi2_terms)

if constraints:
constraints = z.reduce_sum(
[c.value() - log_offset_val for c in constraints]
)
chi2_term += constraints

return chi2_term

@property
def is_extended(self):
return False

def _get_params(
self,
floating: bool | None = True,
is_yield: bool | None = None,
extract_independent: bool | None = True,
) -> set[ZfitParameter]:
if not self.is_extended:
is_yield = False  # the loss does not depend on the yields
return super()._get_params(floating, is_yield, extract_independent)

[docs]
class ExtendedBinnedChi2(BaseBinned):
def __init__(
self,
model: ztyping.BinnedPDFInputType,
data: ztyping.BinnedDataInputType,
constraints: ConstraintsInputType = None,
options: OptionsInputType = None,
):
r"""Binned Chi2 loss, using the :math:N_{tot} from the PDF.

.. math::
\chi^2 = \sum_{\mathrm{bins}} \left( \frac{N_\mathrm{PDF,bin} - N_\mathrm{Data,bin}}{\sigma_\mathrm{Data,bin}} \right)^2

where

.. math::
N_\mathrm{PDF,bin} = \mathrm{pdf}(\text{integral}) \cdot N_\mathrm{PDF,expected}
\sigma_\mathrm{bin} = \text{variance}

with variance the value of :class:~zfit.data.BinnedData.variances of the binned data.

|@doc:loss.init.binned.explain.chi2zeros| If the dataset has empty bins, the errors
will be zero and :math:\chi^2 is undefined. Two possibilities are available and
can be given as an option:

- "empty": "ignore" will ignore all bins with zero entries and won't count to the loss
- "errors": "expected" will use the expected counts from the model
with a Poissonian uncertainty |@docend:loss.init.binned.explain.chi2zeros|

Args:
model: |@doc:loss.binned.init.model| Binned PDF(s) that return the normalized probability
(rel_counts or counts) for
*data* under the given parameters.
If multiple model and data are given, they will be used
in the same order to do a simultaneous fit. |@docend:loss.binned.init.model|
data: |@doc:loss.binned.init.data| Binned dataset that will be given to the *model*.
If multiple model and data are given, they will be used
in the same order to do a simultaneous fit. |@docend:loss.binned.init.data|
constraints: |@doc:loss.init.constraints| Auxiliary measurements ("constraints")
that add a likelihood term to the loss.

.. math::
\mathcal{L}(\theta) = \mathcal{L}_{unconstrained} \prod_{i} f_{constr_i}(\theta)

Usually, an auxiliary measurement -- by its very nature -S  should only be added once
to the loss. zfit does not automatically deduplicate constraints if they are given
multiple times, leaving the freedom for arbitrary constructs.

Constraints can also be used to restrict the loss by adding any kinds of penalties. |@docend:loss.init.constraints|
options: |@doc:loss.init.options| Additional options (as a dict) for the loss.
Current possibilities include:

- 'subtr_const' (default True): subtract from each points
log probability density a constant that
is approximately equal to the average log probability
density in the very first evaluation before
the summation. This brings the initial loss value closer to 0 and increases,
especially for large datasets, the numerical stability.

The value will be stored ith 'subtr_const_value' and can also be given
directly.

The subtraction should not affect the minimum as the absolute
value of the NLL is meaningless. However,
with this switch on, one cannot directly compare
different likelihoods absolute value as the constant
may differ! Use create_new in order to have a comparable likelihood
between different losses or use the full argument in the value function
to calculate the full loss with all constants.

These settings may extend over time. In order to make sure that a loss is the
same under the same data, make sure to use create_new instead of instantiating
a new loss as the former will automatically overtake any relevant constants
and behavior. |@docend:loss.init.options|
"""
self._errordef = 1.0
if options is None:
options = {}
if options.get("empty") is None:
options["empty"] = "ignore"
if options.get("errors") is None:
options["errors"] = "data"
super().__init__(
model=model, data=data, constraints=constraints, options=options
)

def _precompile(self):
super()._precompile()
ignore_empty = (
self._options.get("empty") == "ignore"
or self._options.get("errors") == "expected"
)
data = self.data
_check_small_counts_chi2(data, ignore_empty)

@z.function(wraps="loss", keepalive=True)
def _loss_func(
self,
model: Iterable[ZfitBinnedPDF],
data: Iterable[ZfitBinnedData],
fit_range,
constraints,
log_offset,
):
del fit_range
ignore_empty = self._options.get("empty") == "ignore"
chi2_terms = []
if log_offset is False:
log_offset_val = 0.0
else:
log_offset_val = log_offset
log_offset_val = znp.asarray(log_offset_val, dtype=znp.float64)
for mod, dat in zip(model, data):
values = dat.values(  # TODO: right order of model and data?
# obs=mod.obs
)
probs = mod.counts(dat)
variance_method = self._options.get("errors")
if variance_method == "expected":
variances = znp.sqrt(probs + znp.asarray(1e-307, dtype=znp.float64))
elif variance_method == "data":
variances = dat.variances()
else:
raise ValueError(f"Variance method {variance_method} not supported")
if variances is None:
raise ValueError(f"variances cannot be None for Chi2: {dat}")

chi2_term = chi2_loss_calc(
probs, values, variances, log_offset, ignore_empty=ignore_empty
)
chi2_terms.append(chi2_term)
chi2_term = znp.sum(chi2_terms)

if constraints:
constraints = z.reduce_sum(
[c.value() - log_offset_val for c in constraints]
)
chi2_term += constraints

return chi2_term

@property
def is_extended(self):
return True