mixle.stats.latent.heterogeneous_mixture module

Finite mixtures whose components may use different distribution families.

HeterogeneousMixtureDistribution with data type T is defined by the density

p_mat(Y) = sum_{k=1}^{K} p_mat(Y|Z=k)*p_mat(Z=k),

where p_mat(Z=k) is a mixture weight and p_mat(Y|Z=k) is the kth component distribution. Component distributions must accept the same observation type T but do not need to share parameters or encoders.

Example: A heterogeneous mixture with weights [0.5, 0.5] and component distribution Exponential(beta) and Gamma(k,theta), has form

p_mat(x_mat) = 0.5*P_0(x; beta) + 0.5*P_1(x; k, theta), for x > 0.0,

where

P_0(x;beta) is an exponential density and P_1(x; k, theta) is a Gamma density.

class HeterogeneousMixtureDistribution(components, w=MISSING, name=None, weights=MISSING)

Bases: SequenceEncodableProbabilityDistribution

Mixture distribution with component-specific observation encoders.

Parameters:

• components (Sequence[SequenceEncodableProbabilityDistribution])

• w (list[float] | ndarray)

• name (str | None)

• weights (list[float] | ndarray)

compute_capabilities()

Return compute-backend metadata shared by heterogeneous mixture components.

compute_declaration()

Return the symbolic declaration for heterogeneous mixture weights and components.

density(x)

Evaluate the density of the heterogeneous mixture at one observation.

See log_density() for details.

Parameters:

x (T) – (T): Single observation from heterogeneous mixture distribution. T is data type of components.

Returns:

Density at x.

Return type:

float

density_semantics()

Return joined density semantics over heterogeneous mixture components.

log_density(x)

Evaluate log-density of heterogeneous mixture distribution at observation x.

A K-component heterogeneous mixture has log-density,

log(p_mat(x)) = log(sum_{z=k}^{K} p_mat(x|z=k)*p_mat(z=k)),

where p_mat(x|z=k) is component-k log-density at x, and p_mat(z=k) = w[k]. A log-sum-exp is used to evaluate the sum inside the log of the right-hand side above. (See mixle.utils.vector.log_sum() for details).

Recall: p_mat(x|z=k) need only be compatible with same data type T. They are need not be the same distribution.

Parameters:

x (T) – (T): Single observation from heterogeneous mixture distribution. T is data type of components.

Returns:

Log-density at x.

Return type:

float

component_log_density(x)

Evaluate component-wise log-density of heterogeneous mixture distribution at observation x.

A K-component heterogeneous mixture has log-density,

log(p_mat(x)) = log(sum_{z=k}^{K} p_mat(x|z=k)*p_mat(z=k)),

where p_mat(x|z=k) is component-k log-density at x, and p_mat(z=k) = w[k].

This function returns an ndarray[float] of length K, containing log(p_mat(x|z=k)) as its k^{th} entry.

Parameters:

x (T) – (T): Single observation from mixture distribution. T is data type of components.

Returns:

Numpy array of floats containing component-wise log-density at x.

Return type:

ndarray

posterior(x)

Obtain the posterior distribution for each heterogeneous mixture component at observation x.

The posterior distribution of component ‘k’ at observation x is given by,

1. p_mat(Z=k|x) = p_mat(x|Z=k)*p_mat(z=k) / p_mat(x),

where

2. p_mat(x) = sum_{k=1}^{K} p_mat(x|Z=k)*p_mat(z=k) = sum_{k=1}^{K} p_mat(x|Z=k)*w[k].

This function returns an ndarray[float] of length K, containing p_mat(Z=k|x) as its k^{th} entry.

Parameters:

x (T) – (T): Single observation from heterogeneous mixture distribution. T is data type of components.

Returns:

Numpy array of floats containing posterior distribution at observation x.

Return type:

ndarray

seq_log_density(x)

Vectorized evaluation of component-wise log-density for encoded sequence x.

Evaluates the log-density of each observation in the encoded sequence x (see log_density() for details).

Arg x must be a Tuple of length two containing and encoded from

HeterogeneousMixtureDataEncoder.seq_encode(data) with data type Sequence[T] for data.

x[0] (List[np.ndarray[int]]): The component ids for each distinct SequenceEncodableProbabilityDistribution

subclass.

x[1] (List[T1,T2,..Tk]): A list of sequence encodings of iid an iid observation sequence for each

‘k’ distinct SequenceEncodableProbabilityDistribution subclasses. The data type for each encoding is assumed to be of type Ti.

The returned value is an ndarray[float] with shape (sz,K), where K is the number of mixture components, and sz is the number of iid observations in the encoded sequence x.

Note: A row-wise log-sum-exp is performed for numerical stability. If a row contains a log-density value of,

-np.inf is returned for the corresponding observation value in the encoded sequence x.

Parameters:

x (tuple[list[ndarray], list[Any]]) – See above for details.

Returns:

Numpy array of floats containing the log_density of each observation in encoded sequence.

Return type:

ndarray

seq_component_log_density(x)

Vectorized evaluation of component-wise log-density for encoded sequence x.

Arg x must be a Tuple of length two containing and encoded from

HeterogeneousMixtureDataEncoder.seq_encode(data) with data type Sequence[T] for data.

x[0] (List[np.ndarray[int]]): The component ids for each distinct SequenceEncodableProbabilityDistribution

subclass.

x[1] (List[T1,T2,..Tk]): A list of sequence encodings of iid an iid observation sequence for each

‘k’ distinct SequenceEncodableProbabilityDistribution subclasses. The data type for each encoding is assumed to be of type Ti.

Creates a 2-d numpy array of floats with vectorized evaluations of component_log_density() stored in the rows corresponding to an observation in encoded sequence x.

The returned value is an ndarray[float] with shape (sz,K), where K is the number of mixture components, and sz is the number of iid observations in the encoded sequence x.

Parameters:

x (tuple[list[ndarray], list[Any]]) – See above for details.

Returns:

2-d numpy array of floats having shape (sz,K), where sz is the number of iid obs in encoded sequence x, and K is the number of mixture components.

Return type:

ndarray

backend_seq_component_log_density(x, engine)

Engine-neutral component log densities for heterogeneous encoded data.

Parameters:

• x (tuple[list[ndarray], list[Any]])

• engine (Any)

Return type:

Any

backend_seq_log_density(x, engine)

Engine-neutral heterogeneous-mixture log-density for encoded data.

Parameters:

• x (tuple[list[ndarray], list[Any]])

• engine (Any)

Return type:

Any

seq_posterior(x)

Vectorized evaluation of posterior of HeterogeneousMixtureDistribution for encoded sequence x.

Arg x must be a Tuple of length two containing and encoded from

HeterogeneousMixtureDataEncoder.seq_encode(data) with data type Sequence[T] for data.

x[0] (List[np.ndarray[int]]): The component ids for each distinct SequenceEncodableProbabilityDistribution

subclass.

x[1] (List[T1,T2,..Tk]): A list of sequence encodings of iid an iid observation sequence for each

‘k’ distinct SequenceEncodableProbabilityDistribution subclasses. The data type for each encoding is assumed to be of type Ti.

Vectorized evaluation the posterior of each observation in the encoded sequence x (see posterior() for details).

The returned value is an ndarray[float] with shape (sz,K), where K is the number of mixture components, and sz is the number of iid observations in the encoded sequence x. Each row contains the posterior of the corresponding encoded observation.

Note: A row-wise log-sum-exp is performed for numerical stability. If a row contains a log-density value of,

-np.inf is returned for the corresponding observation value in the encoded sequence x.

Parameters:

x (tuple[list[ndarray], list[Any]]) – See above for details.

Returns:

Numpy array of floats containing the posterior of each observation in encoded sequence.

Return type:

ndarray

sampler(seed=None)

Return a sampler for this heterogeneous mixture.

Parameters:

seed (Optional[int]) – Seed to set for sampling with RandomState.

Returns:

HeterogeneousMixtureSampler object.

Return type:

HeterogeneousMixtureSampler

estimator(pseudo_count=None)

Return an estimator with one child estimator per component.

Parameters:

pseudo_count (Optional[float]) – Used to inflate sufficient statistics in estimation.

Returns:

HeterogeneousMixtureEstimator object.

Return type:

HeterogeneousMixtureEstimator

decomposition()

Heterogeneous mixture components split along the component axis (logsumexp responsibilities inside a shard; per-component stats SUM-reduce). Components are NOT homogeneous, so there is no stacked-parameter tensor to DTensor-shard (engine_axis=None -> host-shard executor mode).

dist_to_encoder()

Return a data encoder assembled from the component distribution encoders.

Return type:

HeterogeneousMixtureDataEncoder

enumerator()

Returns a HeterogeneousMixtureEnumerator iterating the union of component supports in descending mixture probability order.

Return type:

HeterogeneousMixtureEnumerator

class HeterogeneousMixtureEnumerator(dist)

Bases: DistributionEnumerator

Enumerator over the deduplicated union of heterogeneous component supports.

Parameters:

dist (HeterogeneousMixtureDistribution)

class HeterogeneousMixtureSampler(dist, seed=None)

Bases: DistributionSampler

Sampler that draws a component id and delegates to that component sampler.

Parameters:

• dist (HeterogeneousMixtureDistribution)

• seed (int | None)

sample(size=None, *, batched=True)

Draw iid samples from a heterogeneous mixture distribution.

The data type drawn from ‘comp_samplers’ is type T, corresponding to the data type of the mixture components.

If size is None, a single sample (of data type T) is drawn and returned. If size is not None, ‘size’-iid heterogeneous mixture samples are drawn and returned as a List with data type List[T]. With batched=True (default) component draws are grouped and scattered – bit-identical to the per-draw loop (batched=False) but far faster, since each component sampler owns an independent RNG.

Parameters:

• size (Optional[int]) – Number of iid samples to draw.

• batched (bool) – Vectorize component draws (default); set False for the per-draw loop.

Returns:

Data type T or List[T].

Return type:

Any | list[Any]

class HeterogeneousMixtureAccumulator(accumulators, keys=(None, None))

Bases: SequenceEncodableStatisticAccumulator

EM accumulator for heterogeneous mixture weights and per-component statistics.

Parameters:

• accumulators (list[SequenceEncodableStatisticAccumulator])

• keys (tuple[str | None, str | None])

update(x, weight, estimate)

Update sufficient statistics of HeterogeneousMixtureAccumulator with weighted observation.

Requires previous estimate of HeterogeneousMixtureDistribution.

Weights posterior of ‘estimate’ at x. Adds sum to comp_counts, then passes posterior[i] as weight for x into update() call of accumulator[i].

Parameters:

• x (T) – Observation of heterogeneous mixture distribution.

• weight (float) – Weight for observation.

• estimate (HeterogeneousMixtureDistribution) – Previous iteration of EM estimate for HeterogeneousMixtureDistribution.

Returns:

None.

Return type:

None

initialize(x, weight, rng)

Initialize heterogeneous-mixture sufficient statistics from one weighted observation.

If _init_rng is False, _acc_rng is set with rng. This is done for consistency in initialize and seq_initialize functions.

Initialize heterogeneous mixture weights with a sample from Dirichlet distribution. Each SequenceEncodableStatisticAccumulator is for the mixture components is initialized with a call to accumulator[i].initialize.

Parameters:

• x (T) – Observation of heterogeneous mixture distribution.

• weight (float) – Weight for observation.

• rng (RandomState) – Used to set _acc_rng if not previously set.

Returns:

None.

Return type:

None

seq_initialize(x, weights, rng)

Initialize heterogeneous-mixture sufficient statistics from encoded observations.

If _init_rng is False, _acc_rng is set with rng. This is done for consistency in initialize and seq_initialize functions.

Arg x must be a Tuple of length two containing and encoded from

HeterogeneousMixtureDataEncoder.seq_encode(data) with data type Sequence[T] for data.

x[0] (List[np.ndarray[int]]): The component ids for each distinct SequenceEncodableProbabilityDistribution

subclass.

x[1] (List[T1,T2,..Tk]): A list of sequence encodings of iid an iid observation sequence for each

‘k’ distinct SequenceEncodableProbabilityDistribution subclasses. The data type for each encoding is assumed to be of type Ti.

Vectorized implementation of initialize(), for sequence encoded x.

Parameters:

• x (tuple[list[ndarray], list[Any]]) – See above for details.

• weights (ndarray[float]) – Numpy array of positive valued floats.

• rng (RandomState) – Used to set _acc_rng if not previously set.

Returns:

None.

Return type:

None

seq_update(x, weights, estimate)

Vectorized update of sufficient statistics from encoded sequence of observations x.

Arg x must be a Tuple of length two containing and encoded from

HeterogeneousMixtureDataEncoder.seq_encode(data) with data type Sequence[T] for data.

x[0] (List[np.ndarray[int]]): The component ids for each distinct SequenceEncodableProbabilityDistribution

subclass.

x[1] (List[T1,T2,..Tk]): A list of sequence encodings of iid an iid observation sequence for each

‘k’ distinct SequenceEncodableProbabilityDistribution subclasses. The data type for each encoding is assumed to be of type Ti.

Note: Requires a previous estimate of HeterogeneousMixtureDistribution be passed. This may require seq_initialize() to be invoked prior to performing seq_update() calls.

Seq_update is similar to HeterogeneousMixtureDistribution.seq_posterior(). Results are aggregated to comp_counts and accumulators.

Parameters:

• x (tuple[list[ndarray], list[Any]]) – See above for details.

• weights (np.ndarray[float]) – Numpy array of positive floats.

• estimate (MixtureDistribution) – Previous heterogeneous-mixture estimate from EM.

Returns:

None.

Return type:

None

seq_update_engine(x, weights, estimate, engine)

Engine-resident mixture E-step (numpy or torch).

Component log-densities are scored on the active engine and combined into the posterior responsibilities with an engine softmax; per-component responsibilities are fed to the child accumulators. Mirrors seq_update (heterogeneous components keep per-type encodings).

combine(suff_stat)

Merge the sufficient statistics of suff_stat with HeterogeneousMixtureAccumulator instance.

Arg suff_stat is a Tuple of length two containing,

suff_stat[0] (ndarray[float]): Aggregated component counts, suff_stat[1] (Tuple[T1,…,Tk]): Tuple of K sufficient statistics for the heterogeneous mixture components.

Note: The components k^{th} heterogeneous mixture component is assumed to have sufficient statistics of type Tk.

Parameters:

suff_stat (tuple[ndarray, tuple[Any, ...]]) – See above for details.

Returns:

HeterogeneousMixtureAccumulator object.

Return type:

HeterogeneousMixtureAccumulator

value()

Returns sufficient statistics of MixtureAccumulator instance.

The sufficient statistics value returned (suff_stat) is a Tuple of length two containing,

suff_stat[0] (ndarray[float]): Aggregated component counts, suff_stat[1] (Tuple[T1,…,Tk]): Tuple of K sufficient statistics for the heterogeneous mixture components.

Note: The components k^{th} heterogeneous mixture component is assumed to have sufficient statistics of type Tk.

Returns:

Tuple[np.ndarray[float], Tuple[T1,…,Tk]] as described above.

Return type:

tuple[ndarray, tuple[Any, …]]

from_value(x)

Set sufficient statistics of HeterogeneousMixtureAccumulator instance to x.

The sufficient statistics value ‘x’ is a Tuple of length two containing,

x[0] (ndarray[float]): Aggregated component counts, x[1] (Tuple[T1,…,Tk]): Tuple of K sufficient statistics for the heterogeneous mixture components.

Note: The components k^{th} heterogeneous mixture component is assumed to have sufficient statistics of type Tk.

Parameters:

x (tuple[ndarray, tuple[Any, ...]]) – See above for details.

Returns:

HeterogeneousMixtureAccumulator object.

Return type:

HeterogeneousMixtureAccumulator

key_merge(stats_dict)

Merge this accumulator into keyed sufficient statistics.

Parameters:

stats_dict (Dict[str, Any]) – Dict mapping keys to sufficient statistics.

Returns:

None.

Return type:

None

key_replace(stats_dict)

Replace this accumulator’s statistics from matching keyed values.

Parameters:

stats_dict (Dict[str, Any]) – Dict mapping keys to sufficient statistics.

Returns:

None.

Return type:

None

acc_to_encoder()

Return a data encoder assembled from the component accumulator encoders.

Return type:

HeterogeneousMixtureDataEncoder

class HeterogeneousMixtureAccumulatorFactory(factories, dim, keys=(None, None))

Bases: StatisticAccumulatorFactory

Factory for heterogeneous-mixture accumulators.

Parameters:

• factories (list[StatisticAccumulatorFactory])

• dim (int)

• keys (tuple[str | None, str | None])

make()

Return a new heterogeneous-mixture accumulator.

Return type:

HeterogeneousMixtureAccumulator

class HeterogeneousMixtureEstimator(estimators, fixed_weights=None, suff_stat=None, pseudo_count=None, name=None, keys=(None, None))

Bases: ParameterEstimator

Estimator for heterogeneous mixture weights and component distributions.

Parameters:

• estimators (list[ParameterEstimator])

• fixed_weights (ndarray | None)

• suff_stat (ndarray | None)

• pseudo_count (float | None)

• name (str | None)

• keys (tuple[str | None, str | None])

accumulator_factory()

Return an accumulator factory built from the component estimator factories.

Return type:

HeterogeneousMixtureAccumulatorFactory

estimate(nobs, suff_stat)

Estimate a heterogeneous mixture from aggregated sufficient statistics.

suff_stat is a two-item tuple containing:

suff_stat[0] (np.ndarray): Sufficient statistics for mixture weights. suff_stat[1] (Tuple[T1,…,Tk]): Tuple of K sufficient statistics for the heterogeneous mixture components.

suff_stat[1] is passed to the corresponding component estimator.

If fixed_weights is not None, suff_stat[0] is ignored and the estimated distribution uses the fixed weights.

If pseudo_count is set, observed counts are blended with either the estimator’s prior statistics or a uniform pseudo-count allocation.

If pseudo_count is None, observed component counts determine the weights.

Parameters:

• nobs (Optional[float]) – Not used. Kept for consistency with ParameterEstimator super class.

• suff_stat (tuple[ndarray, tuple[Any, ...]]) – See above for details.

Returns:

Estimated mixture distribution.

Return type:

HeterogeneousMixtureDistribution

class HeterogeneousMixtureDataEncoder(encoders)

Bases: DataSequenceEncoder

Encoder that groups observations by compatible component encoder.

Parameters:

encoders (list[DataSequenceEncoder])

seq_encode(x)

Encode a sequence of iid heterogeneous mixture observations.

Note: The data type for every encoder in the keys of HeterogeneousMixtureDataEncoder attribute self.encoder_dict.keys() is T.

The returned tuple contains:

tag_list (List[ndarray[int]): Heterogeneous mixture component ids for encoded sequences in enc_data list. enc_data (List[S1,…,Sm]): A list of ‘m’ encoded sequences of type Sm, corresponding to component ids

in tag_list.

Parameters:

x (Sequence[T]) – A Sequence of iid observations drawn from a heterogeneous mixture distribution.

Returns:

Tuple[List[ndarray[int], List[S1,…,Sm]] as defined above.

Return type:

tuple[list[ndarray], list[Any]]