Discrete distributions¶

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

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

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

    Returns a column with one variate per input row, in the distribution's element dtype
    (`Float64`, `UInt64` or `Boolean`). Output length follows the surrounding context (frame length
    under `select` / `with_columns`, partition length under `over` / `group_by`), and each row's draw
    is derived from a per-row sub-seed mixed from `seed` and the row's position, so the result is
    independent of Polars 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 (the fast path); with any column-valued
    parameter the name follows the first parameter expression (polars root-name semantics, so
    `.name.*` modifiers keep working).
    """
    if self._scalar_kwargs is not None:
        return register_plugin(
            f"{self._plugin_prefix}_sample_scalar",
            (ROW_INDEX_EXPR,),
            kwargs={"seed": seed, **self._scalar_kwargs},
        ).alias("sample")
    return register_plugin(
        f"{self._plugin_prefix}_sample", (*self._param_exprs, ROW_INDEX_EXPR), kwargs={"seed": seed}
    )

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

    Each row's `size` draws are consecutive values from one per-row random stream keyed by `seed` and the
    row's position, so the result is reproducible for a fixed `seed` and independent of Polars chunking and
    thread scheduling. `samples(size=1)` matches `sample` for the same seed, and growing `size` extends each
    row's array without changing the existing draws.

    A row with a null parameter yields a null array (not an array of null elements), produced natively by
    the plugin via the output's outer validity; an invalid parameterisation raises.

    Naming follows `sample`: `"samples"` with all-constant parameters, the first parameter
    expression's root name otherwise.
    """
    if size <= 0:
        msg = f"size must be a positive integer, got {size}"
        raise ValueError(msg)
    out = self._samples(size=size, seed=seed)
    return out.alias("samples") if self._scalar_kwargs is not None else out

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

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

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

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

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.
    """
    q = as_expr(quantile)
    return propagate_null_and_nan(q, self._ppf(q))

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`.
    """
    q = as_expr(quantile)
    return propagate_null_and_nan(q, self._isf(q))

@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(0.5)

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

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

    Returns a column with one variate per input row, in the distribution's element dtype
    (`Float64`, `UInt64` or `Boolean`). Output length follows the surrounding context (frame length
    under `select` / `with_columns`, partition length under `over` / `group_by`), and each row's draw
    is derived from a per-row sub-seed mixed from `seed` and the row's position, so the result is
    independent of Polars 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 (the fast path); with any column-valued
    parameter the name follows the first parameter expression (polars root-name semantics, so
    `.name.*` modifiers keep working).
    """
    if self._scalar_kwargs is not None:
        return register_plugin(
            f"{self._plugin_prefix}_sample_scalar",
            (ROW_INDEX_EXPR,),
            kwargs={"seed": seed, **self._scalar_kwargs},
        ).alias("sample")
    return register_plugin(
        f"{self._plugin_prefix}_sample", (*self._param_exprs, ROW_INDEX_EXPR), kwargs={"seed": seed}
    )

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

    Each row's `size` draws are consecutive values from one per-row random stream keyed by `seed` and the
    row's position, so the result is reproducible for a fixed `seed` and independent of Polars chunking and
    thread scheduling. `samples(size=1)` matches `sample` for the same seed, and growing `size` extends each
    row's array without changing the existing draws.

    A row with a null parameter yields a null array (not an array of null elements), produced natively by
    the plugin via the output's outer validity; an invalid parameterisation raises.

    Naming follows `sample`: `"samples"` with all-constant parameters, the first parameter
    expression's root name otherwise.
    """
    if size <= 0:
        msg = f"size must be a positive integer, got {size}"
        raise ValueError(msg)
    out = self._samples(size=size, seed=seed)
    return out.alias("samples") if self._scalar_kwargs is not None else out

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

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

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

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

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.
    """
    q = as_expr(quantile)
    return propagate_null_and_nan(q, self._ppf(q))

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`.
    """
    q = as_expr(quantile)
    return propagate_null_and_nan(q, self._isf(q))

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(0.5)

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

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

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

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

def entropy(self) -> pl.Expr:
    """Shannon entropy, ``-p * log(p) - (1 - p) * log1p(-p)``.

    Uses the convention ``0 * log 0 = 0`` so the result is ``0`` at the degenerate endpoints ``p in {0, 1}``.
    The second term goes through ``log1p`` for the reason in ``_log_pmf``.
    """
    p = self._checked_p
    q = 1 - p

    return pl.when((p == 0) | (p == 1)).then(0.0).otherwise(-p * p.log() - q * (-p).log1p())

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

    Returns a column with one variate per input row, in the distribution's element dtype
    (`Float64`, `UInt64` or `Boolean`). Output length follows the surrounding context (frame length
    under `select` / `with_columns`, partition length under `over` / `group_by`), and each row's draw
    is derived from a per-row sub-seed mixed from `seed` and the row's position, so the result is
    independent of Polars 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 (the fast path); with any column-valued
    parameter the name follows the first parameter expression (polars root-name semantics, so
    `.name.*` modifiers keep working).
    """
    if self._scalar_kwargs is not None:
        return register_plugin(
            f"{self._plugin_prefix}_sample_scalar",
            (ROW_INDEX_EXPR,),
            kwargs={"seed": seed, **self._scalar_kwargs},
        ).alias("sample")
    return register_plugin(
        f"{self._plugin_prefix}_sample", (*self._param_exprs, ROW_INDEX_EXPR), kwargs={"seed": seed}
    )

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

    Each row's `size` draws are consecutive values from one per-row random stream keyed by `seed` and the
    row's position, so the result is reproducible for a fixed `seed` and independent of Polars chunking and
    thread scheduling. `samples(size=1)` matches `sample` for the same seed, and growing `size` extends each
    row's array without changing the existing draws.

    A row with a null parameter yields a null array (not an array of null elements), produced natively by
    the plugin via the output's outer validity; an invalid parameterisation raises.

    Naming follows `sample`: `"samples"` with all-constant parameters, the first parameter
    expression's root name otherwise.
    """
    if size <= 0:
        msg = f"size must be a positive integer, got {size}"
        raise ValueError(msg)
    out = self._samples(size=size, seed=seed)
    return out.alias("samples") if self._scalar_kwargs is not None else out

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

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

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

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

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.
    """
    q = as_expr(quantile)
    return propagate_null_and_nan(q, self._ppf(q))

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`.
    """
    q = as_expr(quantile)
    return propagate_null_and_nan(q, self._isf(q))

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(0.5)

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

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

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}``.
    """
    # Unlike ``mean`` / ``variance`` there is no closed form, so the sum is evaluated by ``binomial_entropy``
    # in Rust. For column parameters that runs once per row; for scalar parameters it is computed **once** on
    # length-1 inputs and broadcast to length-n behind the ``_moment`` validity gate, so a constant's support sum is
    # not re-evaluated on every row.
    if self._scalar_kwargs is None:
        return register_plugin("binomial_entropy", (self._n, self._p))
    return self._moment(register_plugin("binomial_entropy", self._scalar_lit_args()))