Source code for pyhs3.distributions.basic
"""
Basic HS3 Distribution implementations.
Provides classes for handling basic probability distributions including
Gaussian, Uniform, Poisson, Exponential, Log-Normal, and Landau distributions
as defined in the HS3 specification.
"""
from __future__ import annotations
import math
from typing import Literal, cast
import pytensor.tensor as pt
from pyhs3.context import Context
from pyhs3.distributions.core import Distribution
from pyhs3.typing.aliases import TensorVar
[docs]
class GaussianDist(Distribution):
r"""
Gaussian (normal) probability distribution.
Implements the standard Gaussian probability density function:
.. math::
f(x; \mu, \sigma) = \frac{1}{\sigma\sqrt{2\pi}} \exp\left(-\frac{(x-\mu)^2}{2\sigma^2}\right)
Parameters:
mean (str): Parameter name for the mean (μ).
sigma (str): Parameter name for the standard deviation (sigma).
x (str): Input variable name.
HS3 Reference:
:ref:`hs3:hs3.gaussian-normal-distribution`
"""
type: Literal["gaussian_dist"] = "gaussian_dist"
mean: str | float | int
sigma: str | float | int
x: str | float | int
def likelihood(self, context: Context) -> TensorVar:
"""
Builds a symbolic expression for the Gaussian PDF.
Args:
context (dict): Mapping of names to pytensor variables.
Returns:
pytensor.tensor.variable.TensorVariable: Symbolic representation of the Gaussian PDF.
"""
# log.info("parameters: ", parameters)
norm_const = 1.0 / (pt.sqrt(2 * math.pi) * context[self._parameters["sigma"]])
exponent = pt.exp(
-0.5
* (
(context[self._parameters["x"]] - context[self._parameters["mean"]])
/ context[self._parameters["sigma"]]
)
** 2
)
return cast(TensorVar, norm_const * exponent)
[docs]
class PoissonDist(Distribution):
r"""
Poisson probability distribution.
Implements the Poisson probability mass function:
.. math::
P(k; \lambda) = \frac{\lambda^k e^{-\lambda}}{k!}
Parameters:
mean (str): Parameter name for the rate parameter (λ).
x (str): Input variable name (discrete count).
HS3 Reference:
:ref:`hs3:hs3.dist:poisson`
"""
type: Literal["poisson_dist"] = "poisson_dist"
mean: str | float | int
x: str | float | int
def likelihood(self, context: Context) -> TensorVar:
"""
Builds a symbolic expression for the Poisson PMF.
Args:
context (dict): Mapping of names to pytensor variables.
Returns:
pytensor.tensor.variable.TensorVariable: Symbolic representation of the Poisson PMF.
"""
mean = context[self._parameters["mean"]]
x = context[self._parameters["x"]]
# Poisson PMF: λ^k * e^(-λ) / k!
# Using pt.gammaln for log(k!) = log(Γ(k+1))
log_pmf = x * pt.log(mean) - mean - pt.gammaln(x + 1)
return cast(TensorVar, pt.exp(log_pmf))
[docs]
class ExponentialDist(Distribution):
r"""
Exponential probability distribution.
Implements the exponential probability density function:
.. math::
f(x; c) = \exp(-c \cdot x)
Parameters:
x (str): Input variable name.
c (str): Rate/decay parameter (coefficient).
Note:
The HS3 specification uses the form exp(-c*x), which matches ROOT's RooExponential
when the negateCoefficient flag is True. ROOT handles parameter transformations
automatically for compatibility.
HS3 Reference:
:hs3:label:`exponential_dist <hs3.exponential-distribution>`
ROOT Reference:
:rootref:`RooExponential <classRooExponential.html>`
"""
type: Literal["exponential_dist"] = "exponential_dist"
x: str | float | int
c: str | float | int
def likelihood(self, context: Context) -> TensorVar:
"""
Builds a symbolic expression for the exponential PDF.
Args:
context (dict): Mapping of names to pytensor variables.
Returns:
pytensor.tensor.variable.TensorVariable: Symbolic representation of exponential PDF.
"""
x = context[self._parameters["x"]]
c = context[self._parameters["c"]]
# Exponential PDF: exp(-c * x)
return cast(TensorVar, pt.exp(-c * x))
[docs]
class LogNormalDist(Distribution):
r"""
Log-normal probability distribution.
Implements the log-normal probability density function:
.. math::
f(x; \mu, \sigma) = \frac{1}{x\sigma\sqrt{2\pi}} \exp\left(-\frac{(\ln(x)-\mu)^2}{2\sigma^2}\right)
Parameters:
x (str): Input variable name (must be > 0).
mu (str): Location parameter (log-scale mean).
sigma (str): Scale parameter (log-scale standard deviation).
Note:
This implementation uses the standard parametrization where mu and sigma
are the mean and standard deviation of the underlying normal distribution
in log-space. ROOT handles parameter transformations automatically for
compatibility with median/shape parametrization.
HS3 Reference:
:hs3:label:`lognormal_dist <hs3.log-normal-distribution>`
"""
type: Literal["lognormal_dist"] = "lognormal_dist"
x: str | float | int
mu: str | float | int
sigma: str | float | int
def likelihood(self, context: Context) -> TensorVar:
"""
Builds a symbolic expression for the log-normal PDF.
Args:
context (dict): Mapping of names to pytensor variables.
Returns:
pytensor.tensor.variable.TensorVariable: Symbolic representation of log-normal PDF.
"""
x = context[self._parameters["x"]]
mu = context[self._parameters["mu"]]
sigma = context[self._parameters["sigma"]]
# Log-normal PDF: (1/x) * exp(-((ln(x) - mu)^2) / (2 * sigma^2))
log_x = pt.log(x)
normalized_log = (log_x - mu) / sigma
return cast(TensorVar, (1.0 / x) * pt.exp(-0.5 * normalized_log**2))
[docs]
class LandauDist(Distribution):
r"""
Landau probability distribution.
Implements the Landau probability density function as defined in ROOT's
RooLandau. Used primarily in high-energy physics for modeling energy
loss distributions.
Approximation using modified Gaussian with asymmetric tails:
.. math::
f(x; \mu, \sigma) = \frac{1}{\sigma} \exp\left(-\frac{1}{2}z^2 - \frac{1}{10}(z-1)^2\right)
where $z = \frac{x-\mu}{\sigma}$ for $z > 1$.
Parameters:
x (str): Input variable name.
mean (str): Location parameter.
sigma (str): Scale parameter.
Note:
The Landau distribution is asymmetric with a long tail towards larger values.
This implementation uses an approximation since the exact Landau function
is not available in PyTensor.
HS3 Reference:
Note: Landau distribution is not explicitly defined in the current HS3 specification.
ROOT Reference:
:rootref:`RooLandau <classRooLandau.html>`
"""
type: Literal["landau_dist"] = "landau_dist"
x: str | float | int
mean: str | float | int
sigma: str | float | int
def likelihood(self, context: Context) -> TensorVar:
"""
Builds a symbolic expression for the Landau PDF.
Args:
context (dict): Mapping of names to pytensor variables.
Returns:
pytensor.tensor.variable.TensorVariable: Symbolic representation of Landau PDF.
Note:
This implementation uses a Gaussian approximation. In practice,
ROOT uses more sophisticated approximations or numerical methods.
"""
x = context[self._parameters["x"]]
mean = context[self._parameters["mean"]]
sigma = context[self._parameters["sigma"]]
# Normalized variable
z = (x - mean) / sigma
# Landau approximation using a modified Gaussian with asymmetric tails
# This is a simplified approximation - ROOT uses more sophisticated methods
gaussian_core = pt.exp(-0.5 * z**2)
asymmetric_factor = pt.exp(-0.1 * pt.maximum(0.0, z - 1) ** 2)
return cast(TensorVar, (1.0 / sigma) * gaussian_core * asymmetric_factor)
# Registry of basic distributions
distributions: dict[str, type[Distribution]] = {
"gaussian_dist": GaussianDist,
"uniform_dist": UniformDist,
"poisson_dist": PoissonDist,
"exponential_dist": ExponentialDist,
"lognormal_dist": LogNormalDist,
"landau_dist": LandauDist,
}
# Define what should be exported from this module
__all__ = [
"ExponentialDist",
"GaussianDist",
"LandauDist",
"LogNormalDist",
"PoissonDist",
"UniformDist",
"distributions",
]