class

Dirichlet

extendsExponentialFamily

Dirichlet(concentration: Tensor, validate_args: bool | None = None)

source

Dirichlet distribution on the $K$ -simplex.

Multivariate generalisation of the lucid.distributions.Beta distribution: a distribution over probability vectors $\mathbf{x} \in \Delta^{K-1}$ with $x_i \geq 0$ and $\sum_i x_i = 1$ . It is the conjugate prior of the lucid.distributions.Categorical / lucid.distributions.Multinomial likelihoods and a foundational building block in topic models (LDA), Bayesian mixture models, and population genetics.

The last dimension of concentration is the simplex/event dimension; all preceding dimensions form the batch shape.

Parameters

concentrationTensor

Concentration vector

\boldsymbol{\alpha}

with all entries

\alpha_i > 0

. Shape (..., K).

validate_argsbool= None

If True, validate parameter constraints at construction time.

Notes

Probability density on the $K$ -simplex ( $x_i > 0,\;\sum_i x_i = 1$ ):

p(\mathbf{x}; \boldsymbol{\alpha}) = \frac{1}{B(\boldsymbol{\alpha})} \prod_{i=1}^{K} x_i^{\alpha_i - 1}, \qquad B(\boldsymbol{\alpha}) = \frac{\prod_i \Gamma(\alpha_i)}{\Gamma(\alpha_0)}, \quad \alpha_0 = \sum_i \alpha_i

Moments (with $\alpha_0 = \sum_i \alpha_i$ , $\mu_i = \alpha_i/\alpha_0$ ):

\mathbb{E}[X_i] = \mu_i, \qquad \mathrm{Var}[X_i] = \frac{\mu_i (1 - \mu_i)}{\alpha_0 + 1}, \qquad \mathrm{Cov}[X_i, X_j] = -\frac{\mu_i \mu_j}{\alpha_0 + 1} \;\;(i \neq j)

Special cases:

$K = 2$ → $\mathrm{Beta}(\alpha_1, \alpha_2)$ (after dropping one redundant coordinate).
$\boldsymbol{\alpha} = \mathbf{1}$ → uniform over the simplex.
$\alpha_0 \to \infty$ with fixed $\boldsymbol{\mu}$ → mass concentrates at $\boldsymbol{\mu}$ .

Conjugacy: observing categorical counts $\mathbf{n} = (n_1, \ldots, n_K)$ updates $\mathrm{Dirichlet}(\boldsymbol{\alpha}) \to \mathrm{Dirichlet}(\boldsymbol{\alpha} + \mathbf{n})$ .

Sampling uses the normalised-Gamma method: draw independent $G_i \sim \mathrm{Gamma}(\alpha_i, 1)$ and set $X_i = G_i / \sum_j G_j$ . Samples are detached.

Examples

>>> import lucid
>>> from lucid.distributions import Dirichlet
>>> d = Dirichlet(lucid.tensor([1.0, 2.0, 3.0]))
>>> d.mean  # α / Σ α
Tensor([0.1667, 0.3333, 0.5000])
>>> d.sample((4,))
Tensor([...])

Methods (6)

dunder

init

→None

__init__(concentration: Tensor, validate_args: bool | None = None)

source

Construct a Dirichlet distribution.

Parameters

concentrationTensor

Concentration vector

\boldsymbol{\alpha}

with all entries

> 0

. The last dimension is the event (simplex) dimension

K

; all preceding dimensions form the batch shape.

validate_argsbool | None= None

If True, validate parameter constraints at construction time.

Notes

The Dirichlet distribution with concentration $\boldsymbol{\alpha}$ has PDF over the $K$ -simplex:

p(\mathbf{x}; \boldsymbol{\alpha}) = \frac{1}{B(\boldsymbol{\alpha})} \prod_{i=1}^{K} x_i^{\alpha_i - 1}

where $B(\boldsymbol{\alpha}) = \prod_i \Gamma(\alpha_i) / \Gamma(\sum_i \alpha_i)$ .

Sampling uses the normalised-Gamma trick: draw independent $G_i \sim \text{Gamma}(\alpha_i, 1)$ then return $\mathbf{x} = \mathbf{G} / \sum_i G_i$ .

Examples

>>> import lucid
>>> from lucid.distributions import Dirichlet
>>> d = Dirichlet(lucid.tensor([1.0, 2.0, 3.0]))
>>> d.mean  # proportional to concentration
Tensor([0.1667, 0.3333, 0.5000])

prop

mean

→Tensor

mean: Tensor

source

Expected value of the Dirichlet distribution.

Each component of the mean equals the normalised concentration:

E[X_i] = \frac{\alpha_i}{\sum_j \alpha_j}

Returns

Tensor

Mean vector on the simplex, shape batch_shape + event_shape.

Examples

>>> Dirichlet(lucid.tensor([2.0, 2.0])).mean
Tensor([0.5, 0.5])

prop

variance

→Tensor

variance: Tensor

source

Variance of the Dirichlet distribution (component-wise).

\operatorname{Var}[X_i] = \frac{\mu_i (1 - \mu_i)}{\alpha_0 + 1}, \quad \alpha_0 = \sum_j \alpha_j,\; \mu_i = \alpha_i / \alpha_0

Returns

Tensor

Variance vector, shape batch_shape + event_shape.

sample

→Tensor

sample(sample_shape: tuple[int, ...] = ())

source

Draw samples from the Dirichlet distribution.

Uses the normalised-Gamma method:

\mathbf{x} = \frac{\mathbf{g}}{\sum_i g_i}, \quad g_i \sim \text{Gamma}(\alpha_i, 1)

The result lies on the probability simplex and is detached.

Parameters

sample_shapetuple[int, ...]= ()

Leading shape dimensions for the sample batch. Default is ().

Returns

Tensor

Simplex-valued samples of shape sample_shape + batch_shape + event_shape.

Examples

>>> d = Dirichlet(lucid.tensor([1.0, 1.0, 1.0]))
>>> x = d.sample((100,))
>>> x.sum(dim=-1)  # all ones

log_prob

→Tensor

log_prob(value: Tensor)

source

Log-density of value under the Dirichlet distribution.

\log p(\mathbf{x}; \boldsymbol{\alpha}) = \sum_i (\alpha_i - 1) \log x_i - \log B(\boldsymbol{\alpha})

Parameters

valueTensor

Simplex-valued observations, last dimension is

K

Returns

Tensor

Log-densities, shape batch_shape.

entropy

→Tensor

entropy()

source

Shannon entropy of the Dirichlet distribution (in nats).

H = \log B(\boldsymbol{\alpha}) + (\alpha_0 - K) \psi(\alpha_0) - \sum_i (\alpha_i - 1) \psi(\alpha_i)

where $\alpha_0 = \sum_i \alpha_i$ , $K$ is the number of categories, and $\psi$ is the digamma function.

Returns

Tensor

Entropy in nats, shape batch_shape.