Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Upgrade PDFs, add labels #75

Merged
merged 2 commits into from
Apr 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions requirements_dev.txt
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ pytest>=3.4.2
pytest-cov
pytest-rerunfailures>=6
pytest-runner>=2.11.1
pytest-xdist
seed_intersphinx_mapping
setupext-janitor
Sphinx>=3.5.4
Expand Down
41 changes: 32 additions & 9 deletions zfit_physics/models/pdf_argus.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,12 +48,15 @@ def argus_func(
class Argus(zfit.pdf.BasePDF):
def __init__(
self,
obs: ztyping.ObsTypeInput,
*,
m0,
c,
p,
name: str = "ArgusPDF",
obs: ztyping.ObsTypeInput,
extended: ztyping.ParamTypeInput | None = None,
norm: ztyping.NormTypeInput = None,
name: str = "ArgusPDF",
label: str | None = None,
):
r"""`ARGUS shape <https://en.wikipedia.org/wiki/ARGUS_distribution>`_ describing the invariant mass of a particle
in a continuous background.
Expand All @@ -71,29 +74,49 @@ def __init__(
The implementation follows the `RooFit version <https://root.cern.ch/doc/master/classRooArgusBG.html>`_

Args:
obs: Observable the PDF is defined on
obs: |@doc:pdf.init.obs| Observables of the
model. This will be used as the default space of the PDF and,
if not given explicitly, as the normalization range.

The default space is used for example in the sample method: if no
sampling limits are given, the default space is used.

The observables are not equal to the domain as it does not restrict or
truncate the model outside this range. |@docend:pdf.init.obs|
m0: Maximal energetically allowed mass, cutoff
c: Shape parameter; "peakiness" of the distribution
p: Generalization of the ARGUS shape, for p = 0.5, the normal ARGUS shape is recovered
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|
norm: |@doc:pdf.init.norm| Normalization of the PDF.
By default, this is the same as the default space of the PDF. |@docend:pdf.init.norm|
name: |@doc:pdf.init.name| Human-readable name
or label of
the PDF for better identification. |@docend:pdf.init.name|
label: |@doc:pdf.init.label| Label of the PDF, if None is given, it will be the name. |@docend:pdf.init.label|

Returns:
`tf.Tensor`: the values matching the (broadcasted) shapes of the input
"""
params = {"m0": m0, "c": c, "p": p}
super().__init__(obs=obs, name=name, params=params, extended=extended)
super().__init__(obs=obs, name=name, params=params, extended=extended, norm=norm, label=label)

_N_OBS = 1

def _unnormalized_pdf(self, x):
@zfit.supports()
def _unnormalized_pdf(self, x, params):
"""
Calculation of ARGUS PDF value
(Docs: https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.argus.html)
"""
m = zfit.z.unstack_x(x)
m = x[0]

m0 = self.params["m0"]
c = self.params["c"]
p = self.params["p"]
m0 = params["m0"]
c = params["c"]
p = params["p"]
return argus_func(m, m0, c, p)


Expand Down
17 changes: 10 additions & 7 deletions zfit_physics/models/pdf_cmsshape.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ def cmsshape_pdf_func(x, m, beta, gamma):
Based on code from `spark_tnp <https://gitlab.cern.ch/cms-muonPOG/spark_tnp/-/blob/Spark3/RooCMSShape.cc>`_ and
`numba-stats <https://github.com/HDembinski/numba-stats/blob/main/src/numba_stats/cmsshape.py>`_.
"""
x = z.unstack_x(x)
half = 0.5
two = 2.0
t1 = tf.math.exp(-gamma * (x - m))
Expand Down Expand Up @@ -93,6 +92,7 @@ def __init__(
extended: ztyping.ExtendedInputType | None = None,
norm: ztyping.NormInputType | None = None,
name: str = "CMSShape",
label: str | None = None,
):
"""CMSShape PDF.

Expand Down Expand Up @@ -130,14 +130,17 @@ def __init__(
or label of
the PDF for better identification.
Has no programmatical functional purpose as identification. |@docend:pdf.init.name|
label: |@doc:pdf.init.label| Label of the PDF, if None is given, it will be the name. |@docend:pdf.init.label|
"""
params = {"m": m, "beta": beta, "gamma": gamma}
super().__init__(obs=obs, params=params, name=name, extended=extended, norm=norm)

def _unnormalized_pdf(self, x: tf.Tensor) -> tf.Tensor:
m = self.params["m"]
beta = self.params["beta"]
gamma = self.params["gamma"]
super().__init__(obs=obs, params=params, name=name, extended=extended, norm=norm, label=label)

@zfit.supports()
def _unnormalized_pdf(self, x: tf.Tensor, params) -> tf.Tensor:
m = params["m"]
beta = params["beta"]
gamma = params["gamma"]
x = x[0]
return cmsshape_pdf_func(x=x, m=m, beta=beta, gamma=gamma)


Expand Down
28 changes: 23 additions & 5 deletions zfit_physics/models/pdf_conv.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@ def __init__(
ndraws: int = 20000,
*,
extended: ztyping.ParamTypeInput | None = None,
norm: ztyping.NormTypeInput = None,
name: str = "Convolution",
label: str | None = None,
experimental_pdf_normalized=False,
):
"""Numerical Convolution pdf of *func* convoluted with *kernel*.
Expand All @@ -31,12 +33,28 @@ def __init__(
kernel (:py:class:`zfit.pdf.BasePDF`): PDF with `pdf` method that takes x acting as the kernel.
Here x is a `Data` with the obs and limits of *limits*.
limits (:py:class:`zfit.Space`): Limits for the numerical integration.
obs (:py:class:`zfit.Space`): Observables of the class
extended: If the PDF should be extended, i.e. a yield.
ndraws (int): Number of draws for the mc integration
name (str): Human readable name of the pdf
obs: |@doc:pdf.init.obs| Observables of the
model. This will be used as the default space of the PDF and,
if not given explicitly, as the normalization range.

The default space is used for example in the sample method: if no
sampling limits are given, the default space is used.

The observables are not equal to the domain as it does not restrict or
truncate the model outside this range. |@docend:pdf.init.obs|
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|
norm: |@doc:pdf.init.norm| Normalization of the PDF.
By default, this is the same as the default space of the PDF. |@docend:pdf.init.norm|
name: |@doc:pdf.init.name| Human-readable name
or label of
the PDF for better identification. |@docend:pdf.init.name|
label: |@doc:pdf.init.label| Label of the PDF, if None is given, it will be the name. |@docend:pdf.init.label|
"""
super().__init__(obs=obs, pdfs=[func, kernel], params={}, name=name, extended=extended)
super().__init__(obs=obs, pdfs=[func, kernel], params={}, name=name, extended=extended, norm=norm, label=label)
limits = self._check_input_limits(limits=limits)
if limits.n_limits == 0:
msg = "obs have to have limits to define where to integrate over."
Expand Down
15 changes: 8 additions & 7 deletions zfit_physics/models/pdf_cruijff.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,11 +97,12 @@ def __init__(
params = {"mu": mu, "sigmal": sigmal, "alphal": alphal, "sigmar": sigmar, "alphar": alphar}
super().__init__(obs=obs, params=params, extended=extended, norm=norm, name=name, label=label)

def _unnormalized_pdf(self, x):
mu = self.params["mu"]
sigmal = self.params["sigmal"].value()
alphal = self.params["alphal"].value()
sigmar = self.params["sigmar"].value()
alphar = self.params["alphar"].value()
x = z.unstack_x(x)
@zfit.supports()
def _unnormalized_pdf(self, x, params):
mu = params["mu"]
sigmal = params["sigmal"]
alphal = params["alphal"]
sigmar = params["sigmar"]
alphar = params["alphar"]
x = x[0]
return cruijff_pdf_func(x=x, mu=mu, sigmal=sigmal, alphal=alphal, sigmar=sigmar, alphar=alphar)
13 changes: 7 additions & 6 deletions zfit_physics/models/pdf_erfexp.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,10 +92,11 @@ def __init__(
params = {"mu": mu, "beta": beta, "gamma": gamma, "n": n}
super().__init__(obs=obs, params=params, extended=extended, norm=norm, name=name, label=label)

def _unnormalized_pdf(self, x):
mu = self.params["mu"]
beta = self.params["beta"]
gamma = self.params["gamma"]
n = self.params["n"]
x = z.unstack_x(x)
@zfit.supports()
def _unnormalized_pdf(self, x, params):
mu = params["mu"]
beta = params["beta"]
gamma = params["gamma"]
n = params["n"]
x = x[0]
return erfexp_pdf_func(x=x, mu=mu, beta=beta, gamma=gamma, n=n)
28 changes: 25 additions & 3 deletions zfit_physics/models/pdf_kde.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,17 +18,37 @@ def __init__(
data: tf.Tensor,
bandwidth: ztyping.ParamTypeInput,
obs: ztyping.ObsTypeInput,
name: str = "GaussianKDE",
*,
extended: ztyping.ParamTypeInput | None = None,
norm: ztyping.NormTypeInput = None,
name: str = "GaussianKDE",
label: str | None = None,
):
"""Gaussian Kernel Density Estimation using Silverman's rule of thumb.

Args:
data: Data points to build a kernel around
bandwidth: sigmas for the covariance matrix of the multivariate gaussian
obs:
name: Name of the PDF
obs: |@doc:pdf.init.obs| Observables of the
model. This will be used as the default space of the PDF and,
if not given explicitly, as the normalization range.

The default space is used for example in the sample method: if no
sampling limits are given, the default space is used.

The observables are not equal to the domain as it does not restrict or
truncate the model outside this range. |@docend:pdf.init.obs|
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|
norm: |@doc:pdf.init.norm| Normalization of the PDF.
By default, this is the same as the default space of the PDF. |@docend:pdf.init.norm|
name: |@doc:pdf.init.name| Human-readable name
or label of
the PDF for better identification. |@docend:pdf.init.name|
label: |@doc:pdf.init.label| Label of the PDF, if None is given, it will be the name. |@docend:pdf.init.label|
"""
dtype = zfit.settings.ztypes.float
if isinstance(data, zfit.core.interfaces.ZfitData):
Expand Down Expand Up @@ -80,6 +100,8 @@ def dist_kwargs():
obs=obs,
name=name,
extended=extended,
norm=norm,
label=label,
)

# @zfit.supports()
Expand Down
27 changes: 16 additions & 11 deletions zfit_physics/models/pdf_novosibirsk.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@ def novosibirsk_pdf(x, mu, sigma, lambd):
peak = mu
width = sigma
tail = lambd
x = z.unstack_x(x)

cond1 = znp.less(znp.abs(tail), 1e-7)
arg = 1.0 - (x - peak) * tail / width
Expand Down Expand Up @@ -126,6 +125,7 @@ def __init__(
extended: ztyping.ExtendedInputType | None = False,
norm: ztyping.NormInputType | None = None,
name: str = "Novosibirsk",
label: str | None = None,
):
"""Novosibirsk PDF.

Expand Down Expand Up @@ -158,18 +158,23 @@ def __init__(
``ext_*`` methods and the ``counts`` (for binned PDFs). |@docend:pdf.init.extended|
norm: |@doc:pdf.init.norm| Normalization of the PDF.
By default, this is the same as the default space of the PDF. |@docend:pdf.init.norm|
name: |@doc:pdf.init.name| Human-readable name
or label of
the PDF for better identification.
Has no programmatical functional purpose as identification. |@docend:pdf.init.name|
name: |@doc:pdf.init.name| Name of the PDF.
Maybe has implications on the serialization and deserialization of the PDF.
For a human-readable name, use the label. |@docend:pdf.init.name|
label: |@doc:pdf.init.label| Human-readable name
or label of
the PDF for a better description, to be used with plots etc.
Has no programmatical functional purpose as identification. |@docend:pdf.init.label|
"""
params = {"mu": mu, "sigma": sigma, "lambd": lambd}
super().__init__(obs=obs, params=params, name=name, extended=extended, norm=norm)

def _unnormalized_pdf(self, x):
mu = self.params["mu"]
sigma = self.params["sigma"]
lambd = self.params["lambd"]
super().__init__(obs=obs, params=params, name=name, extended=extended, norm=norm, label=label)

@zfit.supports()
def _unnormalized_pdf(self, x, params):
mu = params["mu"]
sigma = params["sigma"]
lambd = params["lambd"]
x = x[0]
return novosibirsk_pdf(x, mu=mu, sigma=sigma, lambd=lambd)


Expand Down
39 changes: 34 additions & 5 deletions zfit_physics/models/pdf_relbw.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@ def relbw_pdf_func(x, m, gamma):
Notes:
Based on code from this [github gist](https://gist.github.com/andrewfowlie/cd0ed7e6c96f7c9e88f85eb3b9665b97#file-bw-py-L87-L110)
"""
x = z.unstack_x(x)
alpha = gamma / m
gamma2 = m**2 * (1.0 + alpha**2) ** 0.5
k = 2.0 ** (3.0 / 2.0) * m**2 * alpha * gamma2 / (np.pi * (m**2 + gamma2) ** 0.5)
Expand All @@ -39,8 +38,11 @@ def __init__(
m: ztyping.ParamTypeInput,
gamma: ztyping.ParamTypeInput,
obs: ztyping.ObsTypeInput,
name: str = "RelativisticBreitWigner",
*,
extended: ztyping.ParamTypeInput | None = None,
norm: ztyping.NormTypeInput | None = None,
name: str = "RelativisticBreitWigner",
label: str | None = None,
):
"""Relativistic Breit-Wigner distribution.

Expand All @@ -49,11 +51,35 @@ def __init__(
Args:
m: the average value
gamma: the width of the distribution
obs: |@doc:pdf.init.obs| Observables of the
model. This will be used as the default space of the PDF and,
if not given explicitly, as the normalization range.

The default space is used for example in the sample method: if no
sampling limits are given, the default space is used.

The observables are not equal to the domain as it does not restrict or
truncate the model outside this range. |@docend:pdf.init.obs|
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|
norm: |@doc:pdf.init.norm| Normalization of the PDF.
By default, this is the same as the default space of the PDF. |@docend:pdf.init.norm|
name: |@doc:pdf.init.name| Name of the PDF.
Maybe has implications on the serialization and deserialization of the PDF.
For a human-readable name, use the label. |@docend:pdf.init.name|
label: |@doc:pdf.init.label| Human-readable name
or label of
the PDF for a better description, to be used with plots etc.
Has no programmatical functional purpose as identification. |@docend:pdf.init.label|
"""
params = {"m": m, "gamma": gamma}
super().__init__(obs=obs, params=params, name=name, extended=extended)
super().__init__(obs=obs, params=params, name=name, extended=extended, norm=norm, label=label)

def _unnormalized_pdf(self, x: tf.Tensor) -> tf.Tensor:
@zfit.supports()
def _unnormalized_pdf(self, x: tf.Tensor, params) -> tf.Tensor:
"""Calculate the PDF at value(s) x.

Args:
Expand All @@ -62,7 +88,10 @@ def _unnormalized_pdf(self, x: tf.Tensor) -> tf.Tensor:
Returns:
`tf.Tensor`: The value(s) of the unnormalized PDF at x.
"""
return relbw_pdf_func(x, m=self.params["m"], gamma=self.params["gamma"])
m = params["m"]
gamma = params["gamma"]
x = x[0]
return relbw_pdf_func(x, m=m, gamma=gamma)


@z.function(wraps="tensor")
Expand Down
Loading
Loading