Discrete distributions¶

def pmf(self, value: float | IntoExprColumn) -> pl.Expr:
    """Probability mass function, `P(X = value)`. Nulls and NaNs in `value` are propagated."""
    return self._pmf(as_expr(value))

def log_pmf(self, value: float | IntoExprColumn) -> pl.Expr:
    """Natural logarithm of the pmf. Nulls and NaNs in `value` are propagated."""
    return self._log_pmf(as_expr(value))

def sample(self, seed: int | None = None) -> pl.Expr:
    """Draw one random variate per row.

    Output length follows the surrounding context (frame length under `select` / `with_columns`,
    partition length under `over` / `group_by`). Each row's draw derives from a sub-seed mixed from
    `seed` and the row's position, so the result is independent of chunking and thread scheduling.

    A row with an invalid parameter raises; a row with a null parameter yields null. The output is
    named `"sample"` when every parameter is constant; otherwise it follows the first parameter
    expression (polars root-name semantics, so `.name.*` modifiers keep working).
    """
    return self._draw("sample", seed=_checked_seed(seed))

def samples(self, size: int, seed: int | None = None) -> pl.Expr:
    """Draw `size` random variates per row, as `Array(inner=<element dtype>, shape=size)`.

    Each row's draws are consecutive values from one per-row stream keyed by `seed` and the row's
    position, so `samples(size=1)` matches `sample` for the same seed and growing `size` extends
    each row's array without changing the existing draws. A null-parameter row yields a null array
    (outer validity), an invalid parameterisation raises. Named like `sample`, as `"samples"`.
    """
    return self._draw("samples", size=_checked_size(size), seed=_checked_seed(seed))

def cdf(self, value: float | IntoExprColumn) -> pl.Expr:
    """Cumulative distribution function, `P(X <= value)`. Nulls and NaNs in `value` are propagated."""
    return self._cdf(as_expr(value))

def log_cdf(self, value: float | IntoExprColumn) -> pl.Expr:
    """Natural logarithm of the cdf. Nulls and NaNs in `value` are propagated."""
    return self._log_cdf(as_expr(value))

def sf(self, value: float | IntoExprColumn) -> pl.Expr:
    """Survival function, `P(X > value) = 1 - cdf(value)`. Nulls and NaNs in `value` are propagated."""
    return self._sf(as_expr(value))

def log_sf(self, value: float | IntoExprColumn) -> pl.Expr:
    """Natural logarithm of the survival function. Nulls and NaNs in `value` are propagated."""
    return self._log_sf(as_expr(value))

def ppf(self, quantile: float | IntoExprColumn) -> pl.Expr:
    """Percent point function (inverse cdf).

    A `quantile` outside `[0, 1]` yields **null**. Nulls are propagated and a `NaN` quantile yields
    `NaN`, matching scipy.
    """
    return self._ppf(as_expr(quantile))

def isf(self, quantile: float | IntoExprColumn) -> pl.Expr:
    """Inverse survival function, the value `x` with `sf(x) == quantile`.

    Same domain contract as `ppf`, with the endpoints reversed: `quantile` outside `[0, 1]` yields
    null, nulls propagate, `NaN` yields `NaN`.
    """
    return self._isf(as_expr(quantile))

@abstractmethod
def mean(self) -> pl.Expr:
    """Expected value `E[X]`."""

@abstractmethod
def variance(self) -> pl.Expr:
    """Variance `Var[X] = E[(X - E[X])^2]`."""

def std(self) -> pl.Expr:
    """Standard deviation, `sqrt(variance)`."""
    return self.variance().sqrt()

def median(self) -> pl.Expr:
    """Median, `ppf(0.5)`."""
    return self._ppf(as_expr(0.5))

@abstractmethod
def entropy(self) -> pl.Expr:
    """Differential or Shannon entropy, in nats."""

def mean(self) -> pl.Expr:
    """Expected value, ``p``."""
    return self._validated_params

def variance(self) -> pl.Expr:
    """Variance, ``p * (1 - p)``."""
    p = self._raw_p
    return self._moment(p * (1 - p))

def entropy(self) -> pl.Expr:
    """Shannon entropy, ``-p * log(p) - (1 - p) * log1p(-p)``, with ``0 * log 0 = 0`` at ``p in {0, 1}``.

    ``log1p(-p)`` rather than ``log(1 - p)``, which collapses to ``0.0`` below ``p ~ 1.1e-16``.
    """
    p = self._raw_p
    shannon = pl.when((p == 0) | (p == 1)).then(0.0).otherwise(-p * p.log() - (1 - p) * (-p).log1p())
    return self._moment(shannon)

def mean(self) -> pl.Expr:
    """Expected value, ``n * p``."""
    return self._moment(self._n * self._p)

def variance(self) -> pl.Expr:
    """Variance, ``n * p * (1 - p)``."""
    return self._moment(self._n * self._p * (1 - self._p))

def entropy(self) -> pl.Expr:
    """Shannon entropy in nats, the exact support sum ``-sum_k pmf(k) log pmf(k)``.

    ``0`` at the degenerate endpoints ``p in {0, 1}``. There is no closed form, so ``binomial_entropy``
    evaluates the sum in Rust.
    """
    return self._param_plugin("entropy")

def mean(self) -> pl.Expr:
    """Expected value, ``(min + max) / 2``."""
    return self._moment(self._midpoint)

def variance(self) -> pl.Expr:
    """Variance, ``(N**2 - 1) / 12``."""
    return (self.support_size**2 - 1) * ONE_TWELFTH

def median(self) -> pl.Expr:
    """Median, the midpoint ``(min + max) / 2``, which for an even support size is not a support point.

    **Diverges from scipy**, which reports that support point: ``scipy.stats.randint(low=1, high=7).median()``
    is ``3.0`` against this library's ``3.5``.
    """
    return self._moment(self._midpoint)

def entropy(self) -> pl.Expr:
    """Shannon entropy in nats, ``log(N)``."""
    return self.support_size.log()

def mean(self) -> pl.Expr:
    """Expected value, ``1 / p``."""
    return 1 / self._validated_params

def variance(self) -> pl.Expr:
    """Variance, ``(1 - p) / p**2``."""
    p = self._raw_p
    return self._moment((1 - p) / p**2)

def std(self) -> pl.Expr:
    """Standard deviation, ``sqrt(1 - p) / p``.

    Not ``variance().sqrt()``: squaring and unsquaring ``p`` overflows about 150 decades earlier.
    """
    p = self._raw_p
    return self._moment((1 - p).sqrt() / p)

def entropy(self) -> pl.Expr:
    """Shannon entropy, ``(-(1 - p) * log1p(-p) - p * log(p)) / p``, with ``0 * log 0 = 0`` at ``p = 1``."""
    p = self._raw_p
    q = 1 - p
    shannon = pl.when(p == 1).then(0.0).otherwise((-q * (-p).log1p() - p * p.log()) / p)
    return self._moment(shannon)