Source code for zfit.sample

#  Copyright (c) 2023 zfit

from __future__ import annotations

from collections.abc import Iterable

from zfit import z
from zfit.core.interfaces import ZfitPDF
from zfit.util.exception import NotExtendedPDFError
from zfit.z.random import sample_with_replacement, counts_multinomial
from .util.container import convert_to_container

__all__ = ["poisson", "sample_with_replacement", "counts_multinomial"]


[docs] def poisson(n=None, pdfs: Iterable[ZfitPDF] = None): if n and pdfs: raise ValueError("Cannot specify both, `n`, and `pdfs`, at the same time.") if pdfs: pdfs = convert_to_container(pdfs) not_extended = [pdf.is_extended for pdf in pdfs] if not_extended: raise NotExtendedPDFError( f"The following pdfs are not extended but need to be: {not_extended}" ) if len(pdfs) > 1: raise ValueError("More than one model (currently) not supported.") # single pdf only implementation here yield_ = pdfs[0].get_yield() else: yield_ = n poisson_term = z.random.poisson(lam=yield_) return poisson_term