mixle.inference.em module

Expectation-maximization strategy helpers.

The strategies in this module are deliberately orchestration-level objects: they move encoded data through existing estimators/kernels and never contain distribution-specific likelihood math.

class EMStepResult(model, objective=None, accepted=True, metadata=None)

Bases: object

Result from one EM-family strategy step.

Parameters:

• model (SequenceEncodableProbabilityDistribution)

• objective (float | None)

• accepted (bool)

• metadata (dict | None)

class SampledSufficientStatistics(suff_stat, nobs=None)

Bases: object

Unambiguous Monte-Carlo EM payload for an optional observation count.

Parameters:

• suff_stat (Any)

• nobs (float | None)

class EMStrategy(*args, **kwargs)

Bases: Protocol

Structural contract for an EM-family strategy consumed by run_em().

Every strategy object in this module (StandardEM, PosteriorTransformEM, AnnealedEM, …) satisfies this Protocol structurally by exposing a step(...) -> EMStepResult method. run_em and _em_step_fn dispatch on it polymorphically; membership is decided by isinstance().

step(enc_data, estimator, model, engine=..., objective=...)

Run one EM-style update and return the resulting model and objective metadata.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• model (SequenceEncodableProbabilityDistribution)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

EMStepResult

class StandardEM

Bases: object

The ordinary Dempster-Laird-Rubin EM update with an exact M-step.

step(enc_data, estimator, model, engine=None, objective=None)

Run one exact EM update and return the new model.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• model (SequenceEncodableProbabilityDistribution)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

EMStepResult

class CompiledEM(*, compute_dtype=None)

Bases: object

Local full-mixture EM using the profiled fused component kernels.

This is the reusable form of the full-tree execution path used by block EM. It keeps the ordinary EM fixed point and update semantics, but can avoid generic accumulator traversal for fusible subtrees of a MixtureDistribution. Unsupported models or remote execution engines fall back to StandardEM.

Parameters:

compute_dtype (Any)

step(enc_data, estimator, model, engine=None, objective=None)

Run one compiled full-mixture E/M sweep.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• model (SequenceEncodableProbabilityDistribution)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

EMStepResult

class PosteriorTransformEM(temperature=1.0, hard=False)

Bases: object

EM update that transforms mixture posteriors before the M-step.

temperature=1 gives the usual soft EM responsibilities. hard=True gives classification/hard EM. Intermediate temperatures implement a simple deterministic-annealing style generalized EM update.

Parameters:

• temperature (float)

• hard (bool)

step(enc_data, estimator, model, engine=None, objective=None)

Run one posterior-transformed E-step followed by the estimator M-step.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• model (SequenceEncodableProbabilityDistribution)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

EMStepResult

class HardEM

Bases: PosteriorTransformEM

Classification EM using maximum-posterior component assignments.

class AnnealedEM(temperatures, hard_final=False)

Bases: object

Deterministic-annealing EM over a temperature schedule.

Temperatures greater than one flatten mixture responsibilities early in a run, then later entries in the schedule can cool toward ordinary EM at temperature one or hard/classification EM at temperature zero. The object owns only the schedule; posterior math and M-steps remain delegated to PosteriorTransformEM and the estimator.

Parameters:

• temperatures (Sequence[float])

• hard_final (bool)

property current_temperature: float

Return the schedule temperature for the next annealed step.

step(enc_data, estimator, model, engine=None, objective=None)

Run one annealed posterior-transform EM step and advance the schedule.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• model (SequenceEncodableProbabilityDistribution)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

EMStepResult

reset()

Restart the annealing schedule for a new EM run.

Return type:

None

class GeneralizedEM(candidate_fn, require_improvement=True)

Bases: object

Generalized EM wrapper around a caller-supplied candidate step.

The candidate function is called as candidate_fn(enc_data, estimator, model, engine). When require_improvement is true, the candidate is accepted only if the supplied objective (or observed log likelihood by default) does not decrease.

Parameters:

• candidate_fn (Callable[[Any, ParameterEstimator, Any, Any | None], Any])

• require_improvement (bool)

step(enc_data, estimator, model, engine=None, objective=None)

Evaluate and optionally objective-gate one caller-supplied GEM step.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• model (SequenceEncodableProbabilityDistribution)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

EMStepResult

class MonotonicEM(base_strategy=None, require_improvement=True, tolerance=1.0e-9)

Bases: object

Objective-gated wrapper that rejects log-likelihood-decreasing or non-finite steps.

Wraps any base EM-family strategy (StandardEM by default). After the base step it evaluates the objective on the candidate; if the candidate objective is non-finite, or (with require_improvement) it decreases beyond tolerance, the previous model is kept and the step is marked rejected. This is the robust-path guard against the singular-covariance / NaN cascade and against EM steps that overshoot.

Parameters:

• base_strategy (Any | None)

• require_improvement (bool)

• tolerance (float)

step(enc_data, estimator, model, engine=None, objective=None)

Run the base step, then reject it if the objective is non-finite or decreases.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• model (SequenceEncodableProbabilityDistribution)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

EMStepResult

class ConditionalMaximizationEM(conditional_steps, require_improvement=True)

Bases: object

Expectation/conditional-maximization over caller-supplied CM steps.

Parameters:

• conditional_steps (Sequence[Callable[[Any, ParameterEstimator, Any, Any | None], Any]])

• require_improvement (bool)

step(enc_data, estimator, model, engine=None, objective=None)

Run each conditional maximization step with optional objective gates.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• model (SequenceEncodableProbabilityDistribution)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

EMStepResult

class MonteCarloEM(sample_suff_stat_fn, num_samples=1, seed=None)

Bases: object

Monte-Carlo EM over sampled sufficient statistics.

sample_suff_stat_fn is called as fn(enc_data, estimator, model, rng, num_samples, engine). Return a bare sufficient statistic, or SampledSufficientStatistics when the estimator also needs an explicit observation count. A two-tuple is always treated as a statistic because tuple-valued statistics are common.

Parameters:

• sample_suff_stat_fn (Callable[[Any, ParameterEstimator, Any, np.random.RandomState, int, Any | None], Any])

• num_samples (int)

• seed (int | None)

step(enc_data, estimator, model, engine=None, objective=None)

Estimate sufficient statistics by sampling latent completions.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• model (SequenceEncodableProbabilityDistribution)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

EMStepResult

class VariationalEM(variational_step_fn, m_step_fn, initial_state=None, free_energy_fn=None)

Bases: object

Free-energy EM over an explicit variational state.

variational_step_fn updates or creates the variational state. The m_step_fn maps that state to a new model. A supplied free_energy_fn can report the model/state objective without requiring the generic observed-likelihood objective to know about the variational state.

Parameters:

• variational_step_fn (Callable[[Any, ParameterEstimator, Any, Any, Any | None], Any])

• m_step_fn (Callable[[Any, ParameterEstimator, Any, Any, Any | None], Any])

• initial_state (Any)

• free_energy_fn (Callable[[Any, ParameterEstimator, Any, Any, Any | None], float] | None)

step(enc_data, estimator, model, engine=None, objective=None)

Update the variational state, then map it to a candidate model.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• model (SequenceEncodableProbabilityDistribution)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

EMStepResult

class OnlineEM(schedule=None, init_estimator=None, init_p=0.1, rng=None, encoder=None, num_chunks=1, dataset_size=None)

Bases: object

Decay-mode stochastic/online EM over encoded mini-batches.

This adapter exposes StreamingEstimator through the strategy interface used by run_em: each step folds one batch into decayed sufficient statistics and then reuses the estimator’s ordinary M-step.

Parameters:

• schedule (Callable[[int], float] | None)

• init_estimator (ParameterEstimator | None)

• init_p (float)

• rng (np.random.RandomState | None)

• encoder (Any | None)

• num_chunks (int)

• dataset_size (float | None)

step(enc_data, estimator, model, engine=None, objective=None)

Fold one mini-batch into decayed sufficient statistics.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• model (SequenceEncodableProbabilityDistribution)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

EMStepResult

reset()

Drop running statistics before a new online EM run.

Return type:

None

class IncrementalEM(chunk_id_fn=None, init_estimator=None, init_p=0.1, rng=None, encoder=None, num_chunks=1)

Bases: object

Neal-Hinton style incremental EM over replaceable encoded chunks.

Revisited chunks replace their previous sufficient-statistic contribution, allowing repeated passes over partitioned data without re-accumulating the whole dataset each iteration.

Parameters:

• chunk_id_fn (Callable[[Any, ParameterEstimator, Any, Any | None], Any] | None)

• init_estimator (ParameterEstimator | None)

• init_p (float)

• rng (np.random.RandomState | None)

• encoder (Any | None)

• num_chunks (int)

step(enc_data, estimator, model, engine=None, objective=None)

Replace the chunk chosen by chunk_id_fn and update the model.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• model (SequenceEncodableProbabilityDistribution)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

EMStepResult

step_chunk(chunk_id, enc_data, estimator, model, engine=None, objective=None)

Replace one named chunk’s sufficient statistics and update the model.

Parameters:

• chunk_id (Any)

• enc_data (Any)

• estimator (ParameterEstimator)

• model (SequenceEncodableProbabilityDistribution)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

EMStepResult

chunk_value(chunk_id)

Return a stored chunk sufficient-statistic payload.

Parameters:

chunk_id (Any)

Return type:

Any

reset()

Drop stored chunks and running statistics before a new incremental EM run.

Return type:

None

class AcceleratedEM(proposal_fn, base_strategy=None, step_factors=(1.0, 0.5, 0.25), require_improvement=True, tolerance=1.0e-12)

Bases: object

Objective-gated acceleration wrapper around an EM-family strategy.

The wrapped base_strategy performs the ordinary EM/GEM step. The caller-supplied proposal_fn may then propose extrapolated candidates from (old_model, base_model, step_factor, enc_data, estimator, engine). This class owns only the orchestration and objective gate; model-specific extrapolation stays with the caller/model layer.

Parameters:

• proposal_fn (Callable[[Any, Any, float, Any, ParameterEstimator, Any | None], Any])

• base_strategy (Any | None)

• step_factors (Sequence[float])

• require_improvement (bool)

• tolerance (float)

step(enc_data, estimator, model, engine=None, objective=None)

Run the base strategy, test extrapolated candidates, and keep the best.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• model (SequenceEncodableProbabilityDistribution)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

EMStepResult

class SquaremEM(base_strategy=None, packer=None, tolerance=1.0e-12, step_growth=4.0)

Bases: object

SQUAREM acceleration (Varadhan & Roland 2008, SqS3) with an objective gate.

One step runs one SQUAREM cycle – THREE base-strategy sweeps: two ordinary EM sweeps give the secant pair r = theta1 - theta0, v = (theta2 - theta1) - r; the squared-extrapolation proposal theta0 - 2*alpha*r + alpha^2*v with alpha = min(-1, -sqrt(r.r/v.v)) is then stabilized by a third EM sweep. The gate keeps the stabilized proposal only if its objective is at least the plain two-sweep result’s; otherwise the cycle falls back to that plain result, so the accepted sequence is monotone whenever the base strategy is (StandardEM’s exact M-step is). An unpack that rejects the proposal (a constraint violation surfacing as ValueError / FloatingPointError / OverflowError) counts as a rejected cycle, never a crash.

Cost accounting is honest: max_its iterations of optimize(strategy=SquaremEM()) spend up to 3 * max_its E-step sweeps. The probe receipt for the win (overlapping 6-component GMM, n=100k): plain EM needed 200 sweeps to a target log-likelihood that SQUAREM reached in 33 sweeps, every cycle accepted, monotone throughout.

packer: (pack, unpack) between the model and an unconstrained parameter vector; default squarem_packer() (supported families documented there).

The strategy is STATEFUL across steps (the adaptive trust-region cap on alpha): use a fresh instance per fit, exactly as optimize(strategy=SquaremEM()) constructs one.

Parameters:

• base_strategy (Any | None)

• packer (tuple[Callable[[Any], np.ndarray], Callable[[np.ndarray], Any]] | None)

• tolerance (float)

• step_growth (float)

step(enc_data, estimator, model, engine=None, objective=None)

One SQUAREM cycle: two base sweeps, one gated extrapolation sweep.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• model (SequenceEncodableProbabilityDistribution)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

EMStepResult

class RestartEM(initial_models, strategy=None, max_its=10, delta=1.0e-9, max_iter=None)

Bases: object

Run an EM-family strategy from several initial models and keep the best.

Parameters:

• initial_models (Sequence[SequenceEncodableProbabilityDistribution])

• strategy (Any | None)

• max_its (int)

• delta (float | None)

• max_iter (int | None)

run(enc_data, estimator, engine=None, objective=None)

Run each initial model through EM and return the best final model.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• engine (Any | None)

• objective (Callable[[Any], float] | None)

Return type:

SequenceEncodableProbabilityDistribution

run_em(enc_data, estimator, initial_model, strategy=None, max_its=10, delta=1.0e-9, engine=None, objective=None, max_iter=None, monotone=True, tolerance=1.0e-12)

Run an EM-family strategy until convergence or max_its.

objective takes the same values as optimize(): None (MLE), a selection string ('auto' / 'mle' / 'map' / 'vb'), or a ready model -> float callable. max_its is the canonical iteration-cap spelling (matching optimize / fit / best_of); max_iter is accepted as a back-compat alias and overrides max_its when given. monotone=True transactionally rejects an objective decrease independently of whether delta is disabled. Set it to False only for a strategy whose changing surrogate objective intentionally permits temporary decreases in the supplied reporting objective.

The returned model carries a FitProvenance receipt (model.fit_provenance()): iterations actually run, whether the run converged on delta or stopped at the max_its cap (or on a rejected step), and the final objective. A run exiting on the cap reports converged=False, so a truncated fit is distinguishable from a finished one. A run whose very first step is rejected accepts nothing and returns initial_model unchanged; that case additionally raises a UserWarning, because the returned model is the initialization rather than a fit and the receipt alone does not say so.

Parameters:

• enc_data (Any)

• estimator (ParameterEstimator)

• initial_model (SequenceEncodableProbabilityDistribution)

• strategy (EMStrategy | None)

• max_its (int)

• delta (float | None)

• engine (Any | None)

• objective (str | Callable[[Any], float] | None)

• max_iter (int | None)

• monotone (bool)

• tolerance (float)

Return type:

SequenceEncodableProbabilityDistribution

observed_log_likelihood(enc_data, engine=None)

Return a model objective over fixed encoded data.

Parameters:

• enc_data (Any)

• engine (Any | None)

Return type:

Callable[[Any], float]