mixle.stats.bayes.pitman_yor module

Create, estimate, and sample from a Pitman-Yor process distribution over set partitions.

Defines the PitmanYorProcessDistribution, PitmanYorProcessSampler, PitmanYorProcessAccumulatorFactory, PitmanYorProcessAccumulator, PitmanYorProcessEstimator, and the PitmanYorProcessDataEncoder classes for use with mixle.

Data type: List[int] (a partition of n elements given as a cluster-label vector; x[i] is the cluster id of element i, e.g. [0, 0, 1, 0, 2, 1] partitions six elements into blocks of sizes 3, 2, 1). Labels are arbitrary – the distribution is exchangeable and depends only on the block sizes.

The Pitman-Yor process PY(alpha, discount) is the two-parameter generalization of the Dirichlet process (discount = 0 recovers the DP / Chinese Restaurant Process). Its exchangeable partition probability function (EPPF) for a partition of n elements into k blocks of sizes n_1, …, n_k is

p = [prod_{i=1}^{k-1} (alpha + i*discount)] / [(alpha + 1)_{n-1}] * prod_j (1 - discount)_{n_j - 1},

with (x)_m the rising factorial. In log form (via lgamma) this is computed in log_density. Larger alpha and discount favor more blocks; discount controls the heavy tail of the block-size distribution (power-law for discount > 0, exponential for discount = 0).

Sampling uses the sequential “Chinese restaurant” construction over num_elements elements. Estimation fits (alpha, discount) by maximizing the aggregated EPPF log-likelihood; the sufficient statistic is three integer-indexed histograms that capture the (alpha + i)/(alpha + i*discount)/(l - discount) factors exactly across partitions of arbitrary sizes.

class PitmanYorProcessDistribution(alpha=1.0, discount=0.0, num_elements=None, name=None, keys=None)[source]

Bases: SequenceEncodableProbabilityDistribution

Pitman-Yor process over set partitions with concentration alpha and discount in [0, 1).

Data type: List[int] (a cluster-label vector partitioning n elements). discount = 0 is the Dirichlet process / Chinese Restaurant Process.

Parameters:
classmethod compute_capabilities()[source]
density(x)[source]

Return the probability of a partition (cluster-label vector) x.

Parameters:

x (Sequence[int])

Return type:

float

log_density(x)[source]

Return the log-probability of a partition (cluster-label vector) x.

Parameters:

x (Sequence[int])

Return type:

float

seq_log_density(x)[source]

Return vectorized log-probabilities for a sequence of block-size arrays.

Parameters:

x (Sequence[ndarray])

Return type:

ndarray

sampler(seed=None)[source]

Return a sampler for drawing partitions from this distribution.

Parameters:

seed (int | None)

Return type:

PitmanYorProcessSampler

estimator(pseudo_count=None)[source]

Return an estimator that fits alpha (and optionally the discount).

Parameters:

pseudo_count (float | None)

Return type:

PitmanYorProcessEstimator

dist_to_encoder()[source]

Return the data encoder used by this distribution for vectorized methods.

Return type:

PitmanYorProcessDataEncoder

class PitmanYorProcessSampler(dist, seed=None)[source]

Bases: DistributionSampler

Draw iid partitions via the sequential Chinese-restaurant construction.

Parameters:

  • dist (PitmanYorProcessDistribution)

  • seed (int | None)

sample(size=None)[source]

Draw partitions of num_elements elements; a single label vector when size is None.

Parameters:

size (int | None)

Return type:

list[int] | list[list[int]]

class PitmanYorProcessAccumulator(keys=None)[source]

Bases: SequenceEncodableStatisticAccumulator

Accumulate the EPPF histogram sufficient statistics for Pitman-Yor estimation.

Three integer-indexed weighted histograms make the aggregated log-likelihood exact: a_hist[i] counts partitions with n > i (the -log(alpha + i) factors), b_hist[i] counts partitions with k > i (the log(alpha + i*discount) factors), and d_hist[l] counts blocks with size > l (the log(l - discount) factors).

Parameters:

keys (str | None)

update(x, weight, estimate)[source]
Parameters:

  • x (Sequence[int])

  • weight (float)

  • estimate (PitmanYorProcessDistribution | None)

Return type:

None

initialize(x, weight, rng)[source]
Parameters:
Return type:

None

seq_update(x, weights, estimate)[source]
Parameters:
Return type:

None

seq_initialize(x, weights, rng)[source]
Parameters:
Return type:

None

combine(suff_stat)[source]
Parameters:

suff_stat (tuple[float, dict[int, float], dict[int, float], dict[int, float]])

Return type:

PitmanYorProcessAccumulator

value()[source]
Return type:

tuple[float, dict[int, float], dict[int, float], dict[int, float]]

from_value(x)[source]
Parameters:

x (tuple[float, dict[int, float], dict[int, float], dict[int, float]])

Return type:

PitmanYorProcessAccumulator

key_merge(stats_dict)[source]

Pool this accumulator’s statistics into stats_dict under its merge key.

The structural default implements the common single-key pattern: store the accumulator under self.keys the first time the key is seen, else combine into the one already there. Accumulators with several named keys (e.g. an HMM’s init/trans/state keys) or a non-accumulator stats payload override this. A keys of None (the default) is a no-op.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

key_replace(stats_dict)[source]

Replace this accumulator’s statistics from the pooled stats_dict entry (see key_merge).

Parameters:

stats_dict (dict[str, Any])

Return type:

None

acc_to_encoder()[source]
Return type:

PitmanYorProcessDataEncoder

class PitmanYorProcessAccumulatorFactory(keys=None)[source]

Bases: StatisticAccumulatorFactory

Factory for PitmanYorProcessAccumulator.

Parameters:

keys (str | None)

make()[source]
Return type:

PitmanYorProcessAccumulator

class PitmanYorProcessEstimator(discount=0.0, estimate_discount=False, max_alpha=1.0e6, name=None, keys=None)[source]

Bases: ParameterEstimator

Maximum-likelihood estimator for the Pitman-Yor concentration alpha and (optionally) discount.

Parameters:
accumulator_factory()[source]
Return type:

PitmanYorProcessAccumulatorFactory

estimate(nobs, suff_stat)[source]
Parameters:
Return type:

PitmanYorProcessDistribution

class PitmanYorProcessDataEncoder[source]

Bases: DataSequenceEncoder

Encode a sequence of partitions (cluster-label vectors) into per-observation block-size arrays.

seq_encode(x)[source]

Encode the iid observation sequence x for vectorized evaluation.

Parameters:

x (Sequence[Sequence[int]])

Return type:

list[ndarray]