mixle.stats.combinator.select module

Create, estimate, and sample from a select distribution.

Defines the SelectDistribution, SelectSampler, SelectEstimatorAccumulator, SelectEstimatorAccumulatorFactory, SelectEstimator, SelectDataEncoder, and SelectEnumerator classes for use with mixle.

Data type: T (any type accepted by every child distribution). The SelectDistribution routes an observation x to one of its child distributions through a user-supplied choice function c(x) -> {0, …, len(dists)-1}, and evaluates the density of the selected child,

p(x) = p_{c(x)}(x).

The choice function partitions the data space, so each child distribution is estimated only from the observations routed to it.

class TypeDispatch(specs)

Bases: object

Serializable choice function that routes an observation to a child by its Python type.

specs[i] declares the type(s) that child i emits; an observation routes to the FIRST child it is an instance of. Specs are friendly names (‘str’, ‘int’, ‘float’, ‘number’, ‘bytes’, ‘bool’, ‘complex’) or the matching Python type objects, with numpy scalar types handled automatically. Because its state is just those names, it round-trips through JSON with no manual registration – unlike a hand-written lambda choice function. Order matters: put more specific types first (e.g. ‘bool’ before ‘int’, since a bool is also an integer).

Parameters:

specs (Sequence[Any])

class SelectDistribution(dists, choice_function, weights=None)

Bases: SequenceEncodableProbabilityDistribution

SelectDistribution routes each observation to one child distribution via a choice function.

Two modes, set by weights:

• Conditional (weights is None, the default) – the density is p(x) = p_{c(x)}(x): the choice function selects which child scores x and the result is that child’s density. There is no branch distribution, so this is a density conditioned on the branch and cannot be sampled as a standalone generative model.

• Dispatch mixture (weights given) – the density is p(x) = w_{c(x)} * p_{c(x)}(x), a normalized mixture over the union of the child supports in which the choice function plays the role of an observed component label. This is the right model for data of disjoint types (e.g. a mix of strings and numbers, where the type identifies the component): unlike MixtureDistribution the component is observed, so fitting is closed-form (the weights are the branch proportions, each child is fit on its routed subset – no EM) and sampling draws a branch from weights then a value from that child.

For the dispatch-mixture density to be normalized the choice function must partition the observation space consistently with the children, i.e. every value a child can emit must route back to that child (choice_function(x) == k for x in child k’s support).

Parameters:

• dists (Sequence[SequenceEncodableProbabilityDistribution])

• choice_function (Callable[[T], int])

• weights (Sequence[float] | None)

classmethod by_type(children, weights='auto')

Build a dispatch mixture that routes observations to children by their Python type.

The friendly constructor for heterogeneous-typed data: instead of hand-writing (and registering) a choice function, pass each child with the type(s) it emits and the routing is derived automatically via a serializable TypeDispatch.

Example – a mixture of strings and counts:

sel = SelectDistribution.by_type([(str, cat), (int, poisson)])
fitted = estimate(mixed_data, sel.estimator())   # learns branch weights + each child

Parameters:

• children (Sequence[tuple[Any, SequenceEncodableProbabilityDistribution]]) – Sequence of (type_spec, distribution) pairs. type_spec is a type (str, int, float, …), a friendly name ('str', 'number', 'float', 'bytes', 'bool', 'complex'), or a tuple of these. Observations route to the FIRST matching child, so order more specific types first.

• weights (Any) – 'auto' (default) seeds uniform branch weights, giving a fittable generative dispatch mixture whose .estimator() learns the true proportions; None gives the weightless conditional density; or pass an explicit weight sequence.

Returns:

SelectDistribution routing by a TypeDispatch choice function.

Return type:

SelectDistribution

compute_capabilities()

compute_declaration()

density(x)

Density of the child distribution selected for observation x.

In dispatch-mixture mode the selected child’s density is scaled by its branch weight.

Parameters:

x (T) – Observation compatible with the child selected by the choice function.

Returns:

Density of the selected child distribution at x.

Return type:

float

log_density(x)

Log-density of the child distribution selected for observation x.

In dispatch-mixture mode the selected child’s log-density is offset by log(weight).

Parameters:

x (T) – Observation compatible with the child selected by the choice function.

Returns:

Log-density of the selected child distribution at x.

Return type:

float

seq_log_density(x)

Vectorized evaluation of the log-density on sequence encoded data x.

The encoding groups observations by choice index: x[1][i] is the choice index of group i, x[0][i] holds the original positions of the group’s observations, and x[2][i] is the group’s data encoded by the matching child encoder. Each group is scored by the child distribution its choice index selects.

Parameters:

x (tuple[tuple[ndarray, ...], tuple[int, ...], tuple[Any, ...]]) – Sequence encoded data produced by SelectDataEncoder.seq_encode().

Returns:

Numpy array of log-densities, one entry per encoded observation.

Return type:

ndarray

backend_seq_log_density(x, engine)

Engine-neutral vectorized log-density for choice-grouped encodings.

Parameters:

• x (tuple[tuple[ndarray, ...], tuple[int, ...], tuple[Any, ...]])

• engine (Any)

Return type:

Any

classmethod backend_stacked_params(dists, engine)

Return stacked child parameters for homogeneous select-wrapper mixtures.

Parameters:

• dists (Sequence[SelectDistribution])

• engine (Any)

Return type:

dict[str, Any]

classmethod backend_stacked_log_density(x, params, engine)

Return an (n, k) matrix of choice-routed select log densities.

Parameters:

• x (tuple[tuple[ndarray, ...], tuple[int, ...], tuple[Any, ...]])

• params (dict[str, Any])

• engine (Any)

Return type:

Any

classmethod backend_stacked_sufficient_statistics_with_estimator(x, weights, params, engine, estimator)

Return per-component legacy select sufficient statistics.

Parameters:

• x (tuple[tuple[ndarray, ...], tuple[int, ...], tuple[Any, ...]])

• weights (Any)

• params (dict[str, Any])

• engine (Any)

• estimator (Any)

Return type:

tuple[list[tuple[Any, Any]], …]

gradient_fit_state(engine, torch, leaves, recurse, tensor_param)

Return distribution-owned state for autograd fitting.

Parameters:

• engine (Any)

• torch (Any)

• leaves (list[Any])

• recurse (Any)

• tensor_param (Any)

Return type:

Any

to_fisher(**kwargs)

Fisher view for the select combinator.

sampler(seed=None)

Creates a SelectSampler object for sampling from the child distributions.

Parameters:

seed (Optional[int]) – Seed for the random number generator used in sampling.

Returns:

SelectSampler object.

Return type:

SelectSampler

estimator(pseudo_count=None, estimate_weights=None)

Creates a SelectEstimator with one child estimator per child distribution.

Parameters:

• pseudo_count (Optional[float]) – Passed through to each child estimator.

• estimate_weights (Optional[bool]) – Whether the fitted distribution carries branch weights (the dispatch-mixture model). When None (default) this follows the current mode – a weighted distribution re-fits weights, a conditional one stays conditional.

Returns:

SelectEstimator object.

Return type:

SelectEstimator

dist_to_encoder()

Creates a SelectDataEncoder object for encoding sequences of SelectDistribution data.

Returns:

SelectDataEncoder object.

Return type:

SelectDataEncoder

enumerator()

Creates a SelectEnumerator iterating the union of child supports in descending select-density order. All children must support enumeration.

Return type:

SelectEnumerator

class SelectEnumerator(dist)

Bases: DistributionEnumerator

Enumerates the union of the child supports in descending select-density order.

Parameters:

dist (SelectDistribution)

class SelectSampler(dist, seed=None)

Bases: DistributionSampler

SelectSampler draws samples from each child distribution of a SelectDistribution.

Parameters:

• dist (SelectDistribution)

• seed (int | None)

sample(size=None)

Draw from the select distribution.

Dispatch-mixture mode (weights set): a branch is drawn from the weights and a single value is drawn from that child – a true draw from p(x) = w_{c(x)} p_{c(x)}(x).

Conditional mode (weights is None): there is no branch distribution to draw from, so this falls back to sampling each child independently and grouping the draws, returning a tuple with one entry per child (or a list of such tuples when size is given). That is NOT a draw from the select density; provide weights to sample a dispatch mixture.

Parameters:

size (Optional[int]) – Number of draws. If None a single draw is returned.

Returns:

In dispatch-mixture mode a single value (or a list of size values); in conditional mode a tuple with one sample per child (or a list of such tuples).

class SelectEstimatorAccumulator(accumulators, choice_function)

Bases: SequenceEncodableStatisticAccumulator

SelectEstimatorAccumulator accumulates sufficient statistics for each child distribution, routing each observation to one child accumulator via the choice function.

Parameters:

• accumulators (Sequence[SequenceEncodableStatisticAccumulator])

• choice_function (Callable[[T], int])

update(x, weight, estimate)

Route one weighted observation to the accumulator of the selected child.

Parameters:

• x (T) – Observation routed by the choice function.

• weight (float) – Weight for the observation.

• estimate (Optional[SelectDistribution]) – Previous estimate; the matching child distribution is passed through to the child accumulator if provided.

Returns:

None.

Return type:

None

initialize(x, weight, rng)

Initialize the accumulator of the selected child with one weighted observation.

Parameters:

• x (T) – Observation routed by the choice function.

• weight (float) – Weight for the observation.

• rng (RandomState) – RandomState used to seed the per-child RandomStates.

Returns:

None.

Return type:

None

seq_update(x, weights, estimate)

Vectorized update of the child accumulators from sequence encoded data x.

Each encoded group i carries its choice index x[1][i]; the group’s encoded data and weights are routed to the child accumulator at that choice index.

Parameters:

• x (tuple[tuple[ndarray, ...], tuple[int, ...], tuple[Any, ...]]) – Sequence encoded data produced by SelectDataEncoder.seq_encode().

• weights (np.ndarray) – Weights for each encoded observation.

• estimate (Optional[SelectDistribution]) – Previous estimate; the matching child distributions are passed through to the child accumulators if provided.

Returns:

None.

Return type:

None

seq_update_engine(x, weights, estimate, engine)

Engine-resident E-step: each group’s weights are gathered on the active engine and routed to the chosen child accumulator through the engine. Matches seq_update.

Parameters:

• x (tuple[tuple[ndarray, ...], tuple[int, ...], tuple[Any, ...]])

• weights (Any)

• estimate (SelectDistribution | None)

• engine (Any)

Return type:

None

seq_initialize(x, weights, rng)

Vectorized initialization of the child accumulators from sequence encoded data x.

Each encoded group i carries its choice index x[1][i]; the group’s encoded data and weights are routed to the child accumulator at that choice index.

Parameters:

• x (tuple[tuple[ndarray, ...], tuple[int, ...], tuple[Any, ...]]) – Sequence encoded data produced by SelectDataEncoder.seq_encode().

• weights (np.ndarray) – Weights for each encoded observation.

• rng (RandomState) – RandomState used to seed the per-child RandomStates.

Returns:

None.

Return type:

None

combine(suff_stat)

Aggregate sufficient statistics suff_stat with this accumulator’s statistics.

Parameters:

suff_stat (Sequence[Tuple[float, Any]]) – One (weight, child sufficient statistic) pair per child, as returned by value().

Returns:

SelectEstimatorAccumulator with combined sufficient statistics.

Return type:

SelectEstimatorAccumulator

value()

Returns the sufficient statistics as a list of (weight, child value) pairs.

Return type:

list[tuple[float, Any]]

from_value(x)

Set the accumulator’s sufficient statistics to x.

Parameters:

x (Sequence[Tuple[float, Any]]) – One (weight, child sufficient statistic) pair per child, as returned by value().

Returns:

SelectEstimatorAccumulator object.

Return type:

SelectEstimatorAccumulator

scale(c)

Scale linear sufficient statistics in-place by c.

The structural default is correct for ordinary weighted sums, nested tuples/lists/dicts, and numeric arrays. Families whose value() payload includes non-linear metadata such as support bounds must override this method and leave that metadata unscaled.

Parameters:

c (float)

Return type:

SelectEstimatorAccumulator

key_merge(stats_dict)

Invoke key_merge on each child accumulator.

Parameters:

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

Returns:

None.

Return type:

None

key_replace(stats_dict)

Invoke key_replace on each child accumulator.

Parameters:

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

Returns:

None.

Return type:

None

acc_to_encoder()

Creates a SelectDataEncoder object for encoding sequences of SelectDistribution data.

Returns:

SelectDataEncoder object.

Return type:

SelectDataEncoder

class SelectEstimatorAccumulatorFactory(estimators, choice_function)

Bases: StatisticAccumulatorFactory

SelectEstimatorAccumulatorFactory creates SelectEstimatorAccumulator objects from the child estimators.

Parameters:

• estimators (Sequence[ParameterEstimator])

• choice_function (Callable[[T], int])

make()

Creates a SelectEstimatorAccumulator with one child accumulator per child estimator.

Returns:

SelectEstimatorAccumulator object.

Return type:

SelectEstimatorAccumulator

class SelectEstimator(estimators, choice_function, estimate_weights=False)

Bases: ParameterEstimator

SelectEstimator estimates a SelectDistribution from child sufficient statistics.

Parameters:

• estimators (Sequence[ParameterEstimator])

• choice_function (Callable[[T], int])

• estimate_weights (bool)

accumulator_factory()

Creates a SelectEstimatorAccumulatorFactory from the child estimators.

Returns:

SelectEstimatorAccumulatorFactory object.

Return type:

SelectEstimatorAccumulatorFactory

estimate(nobs, suff_stat)

Estimate a SelectDistribution from aggregated sufficient statistics.

Parameters:

• nobs (Optional[float]) – Not used. Kept for consistency.

• suff_stat (Sequence[Tuple[float, Any]]) – One (weight, child sufficient statistic) pair per child, as returned by SelectEstimatorAccumulator.value().

Returns:

SelectDistribution object.

Return type:

SelectDistribution

class SelectDataEncoder(encoders, choice_function)

Bases: DataSequenceEncoder

SelectDataEncoder encodes sequences of SelectDistribution data, grouping observations by their choice index and delegating each group to the matching child encoder.

Parameters:

• encoders (Sequence[DataSequenceEncoder])

• choice_function (Callable[[T], int])

seq_encode(x)

Encode a sequence of iid SelectDistribution observations for vectorized “seq_” calls.

Observations are grouped by their choice index, in order of first appearance. The encoding is a tuple of three aligned tuples (one entry per observed choice index):

rv[0] (Tuple[np.ndarray, …]): Original positions of each group’s observations. rv[1] (Tuple[int, …]): The choice index of each group. rv[2] (Tuple[Any, …]): Each group’s data encoded by the matching child encoder.

Parameters:

x (Sequence[T]) – Sequence of iid observations.

Returns:

See description above.

Return type:

tuple[tuple[ndarray, …], tuple[int, …], tuple[Any, …]]

SelectAccumulator

alias of SelectEstimatorAccumulator

SelectAccumulatorFactory

alias of SelectEstimatorAccumulatorFactory

class SelectFisherView(dist)

Bases: EmpiricalMetricFixedFisherView

Parameters:

dist (Any)