mixle.stats.latent.heterogeneous_mixture module

Create, estimate, and sample from a heterogeneous mixture distribution.

Defines the HeterogeneousMixtureDistribution, HeterogeneousMixtureSampler, HeterogeneousMixtureAccumulatorFactory, HeterogeneousMixtureAccumulator, HeterogeneousMixtureEstimator, and the HeterogeneousMixtureDataEncoder classes for use with mixle.

HeterogeneousMixtureDistribution with data type T, is defined by the density of the form,

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 defined as a the k^{th} component distribution. Note that the component distributions p_mat(Y|Z=k) must only be compatible in data type T.

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

compute_declaration()

density(x)

Evaluate density of heterogeneous mMixture distribution at observation x.

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

What log_density returns relative to the true log-density (default: exact).

Override to declare that this distribution’s log_density is a variational lower bound (ELBO), an upper bound, or an approximation rather than the exact log p(x). This is surfaced as the ExactDensity capability and noted in mixle.describe(), so code that needs an exact likelihood can require(x, ExactDensity) instead of silently trusting a bound.

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)

Create HeterogeneousMixtureSampler for sampling from HeterogeneousMixtureDistribution instance.

Parameters:

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

Returns:

HeterogeneousMixtureSampler object.

Return type:

HeterogeneousMixtureSampler

estimator(pseudo_count=None)

Create HeterogeneousMixtureEstimator for estimating HeterogeneousMixtureDistribution.

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

Returns a HeterogeneousMixtureDataEncoder object for encoding sequences of iid observations from HeterogeneousMixtureDistribution.

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

Parameters:

dist (HeterogeneousMixtureDistribution)

class HeterogeneousMixtureSampler(dist, seed=None)

Bases: DistributionSampler

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

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 HeterogeneousMixtureAccumulator object with weighted observation x.

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)

Vectorized initialization of HeterogeneousMixtureAccumulator object for sequence encoded observations x.

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) – HeterogeneousMixtureDistribution object for previous 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 sufficient statistics of object instance with suff stats containing matching keys.

Parameters:

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

Returns:

None.

Return type:

None

key_replace(stats_dict)

Set sufficient statistics of object instance to suff_stats with matching keys.

Parameters:

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

Returns:

None.

Return type:

None

acc_to_encoder()

Returns a HeterogeneousMixtureDataEncoder object for encoding sequences of iid observations from HeterogeneousMixtureDistribution.

Return type:

HeterogeneousMixtureDataEncoder

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

Bases: StatisticAccumulatorFactory

Parameters:

• factories (list[StatisticAccumulatorFactory])

• dim (int)

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

make()

“Return HeterogeneousMixtureAccumulator object with SequenceEncodableStatisticAccumulator objects for the components and keys passed.

Return type:

HeterogeneousMixtureAccumulator

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

Bases: ParameterEstimator

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

Returns HeterogeneousMixtureAccumulatorFactory object passing component StatisticAccumulatorFactory objects and keys.

Return type:

HeterogeneousMixtureAccumulatorFactory

estimate(nobs, suff_stat)

Estimate HeterogeneousMixtureDistribution from aggregated sufficient statistics.

Args suff_stat is a Tuple length two containing:

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

suff_stat[1] is passed to estimate() function of each corresponding entry in member variable ‘estimators’.

If fixed_weights is not None, suff_stat[0] is not used and the weights of the HeterogeneousMixtureDistribution

are set to fixed_weights.

If pseudo_count is passed, arg suff_stat[0] is aggregated with re-weighted member variable suff_stat. If member variable suff_stat is None, then the arg suff_stat[0] is re-weighted with pseudo_count to estimate the weights.

If pseudo_count is None, ar suff_stat[0] is used to estimate 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:

HeterogeneousMixtureDistribution object.

Return type:

HeterogeneousMixtureDistribution

class HeterogeneousMixtureDataEncoder(encoders)

Bases: DataSequenceEncoder

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.

Returns a Tuple of length two containing:

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]]