mixle.inference.freeze_rollup module

Freeze/roll-up cache – per-datum log-density caching over frozen subtrees (workstream D2).

Frame (see the ConditionalJIT track, D1-D6): the estimator tree is an IR. D1 (mixle.inference.node_report) instruments every node with a per-round residual/Q-gain report and an update_kind classification. D2 spends that report: once a subtree’s D1 report says it has stopped changing (a structurally "frozen" node, or a converged residual/Q-gain), this module skips recomputing its per-datum log-density on every subsequent E-step round and instead reuses a cached array – turning E-step cost from O(full tree) into O(active fraction).

Correctness backbone (unchanged from the rest of the D-track): freeze/roll-up is a SCHEDULING optimization only. It never changes what is computed, only how often – a cached per-datum log-density is byte-identical to a freshly recomputed one for the SAME parameters, and the cache is invalidated (recomputed) the instant a subtree’s parameters move again. The Neal-Hinton free energy F this module tracks is the same observed-data-log-likelihood objective mixle.inference.em.observed_log_likelihood() already tracks for vanilla EM; F is the audit receipt that the cache never silently drifted from a real EM trajectory.

Scope: this module targets mixle.stats.latent.mixture.MixtureDistribution, the combinator with the clearest “some subtrees stop mattering” story (a component whose mixture weight collapses near zero stops moving under EM – its sufficient-statistic contribution is ~0 – and stays frozen once collapsed). The freeze/roll-up mechanics (cache, invalidation, D1-driven detection, active-fraction accounting) generalize to any combinator whose E-step already produces one per-datum log-density array per child (Composite/Sequence), but a single, well-tested combinator is the honest S/M-effort scope for this item; later track items (D3’s scheduler) are expected to widen it.

class FreezeRollupStats(round_index, n_components, n_active, n_frozen, n_zero_weight, n_log_density_evals, objective, accepted=True)[source]

Bases: object

One round’s accounting for the freeze/roll-up cache – the acceptance-criteria receipt.

n_log_density_evals is the “wall-clock proxy” this module actually optimizes: the count of real component.seq_log_density(...) calls issued this round (each O(nobs)), as opposed to a cache hit (O(1), a dict lookup + signature compare). objective is the real, provably-monotone Neal-Hinton F for this round (observed-data log-likelihood), read straight off the E-step this module already had to run – no parallel F-tracking mechanism.

Parameters:

  • round_index (int)

  • n_components (int)

  • n_active (int)

  • n_frozen (int)

  • n_zero_weight (int)

  • n_log_density_evals (int)

  • objective (float)

  • accepted (bool)

property active_fraction: float

Fraction of components genuinely (re-)evaluated this round, out of all components.

class FreezeRollupCache(*, q_gain_tol=_DEFAULT_Q_GAIN_TOL, weight_tol=_DEFAULT_WEIGHT_TOL, weight_delta_tol=_DEFAULT_WEIGHT_DELTA_TOL, freeze_patience=_DEFAULT_FREEZE_PATIENCE)[source]

Bases: object

Per-mixture-component cache of per-datum log-density, keyed by component index.

A component is eligible for caching once mixle.inference.node_report.node_report() reports it as D1-"frozen" (the Neutral capability) or as having a converged residual/ Q-gain (abs(q_gain) < q_gain_tol) sustained for freeze_patience consecutive rounds while its mixture weight has collapsed below weight_tol (the natural mixture-specific trigger: a near-zero-weight component’s data-weighted M-step contribution is ~0, so its params – and hence its residual/Q-gain – stop moving). Once frozen, the caller (see run_em_freeze_rollup()) also skips that component’s M-step entirely and carries its model object forward unchanged, so the cached signature can never go stale on its own; the cache only invalidates if a caller explicitly mutates/re-estimates a “frozen” component (invalidate()) or the component’s own parameter signature is found to have moved (a belt-and-suspenders check on every lookup, not just a documented invariant).

Parameters:
invalidate(idx=None)[source]

Drop a cached entry (idx=None clears the whole cache and freeze streaks).

Call this if a caller explicitly re-triggers a component the cache had frozen (e.g. a later scheduler decides to re-activate it): the next lookup is guaranteed to recompute rather than silently return the stale array.

Parameters:

idx (int | None)

Return type:

None

is_frozen(idx, component, weight)[source]

Return the D1-driven freeze verdict for component idx this round.

Reads a fresh NodeReport every round (cheap: a small Monte-Carlo self-residual, not a real-data pass) so a component that later moves again (unfrozen) is detected immediately – the streak resets to 0 the instant the residual/Q-gain stops looking converged, the weight climbs back out of weight_tol, or the weight itself is still moving.

Both a converged own-residual/Q-gain AND a converged weight are required: a mixture component’s own fit (mean/variance) can stabilize well before the joint E-step’s responsibility reallocation across near-degenerate components finishes settling its weight (slow-manifold EM plateaus). Freezing on residual/Q-gain alone would risk locking in a component’s weight (and hence the M-step never revisiting it) before it has actually reached its coordinate-ascent fixed point – silently changing what the fit converges to, which the ConditionalJIT track’s own correctness backbone forbids (a scheduler may change speed, never the answer). Requiring the weight to also have stopped moving (abs(weight - prev_weight) < weight_delta_tol) for freeze_patience consecutive rounds is the guard against that false-freeze failure mode.

Parameters:
Return type:

bool

component_log_density(idx, component, enc, *, frozen)[source]

Return (log_density, was_cache_hit) for one component on this round.

A cache hit costs a dict lookup + an O(param_count) signature compare – never a call into component.seq_log_density (O(nobs)), which is the entire point of D2. Any mismatch between the cached signature and the component’s current parameters – frozen or not – forces a recompute, so a stale cache can never silently persist past the point where it is wrong (acceptance criterion 3).

Parameters:
Return type:

tuple[ndarray, bool]

detect_frozen(cache, model)[source]

Return the set of component indices D1 reports as frozen for model this round.

Parameters:

  • cache (FreezeRollupCache)

  • model (MixtureDistribution)

Return type:

set[int]

run_em_freeze_rollup(enc_data, estimator, initial_model, *, max_its=10, delta=1.0e-9, cache=None, q_gain_tol=_DEFAULT_Q_GAIN_TOL, weight_tol=_DEFAULT_WEIGHT_TOL, weight_delta_tol=_DEFAULT_WEIGHT_DELTA_TOL, freeze_patience=_DEFAULT_FREEZE_PATIENCE, accept_tolerance=_DEFAULT_ACCEPT_TOLERANCE)[source]

Run EM over a MixtureDistribution with D1-driven freeze/roll-up E-step caching.

Mirrors mixle.inference.em.run_em() + PosteriorTransformEM’s soft-EM update (same GEM/ECM coordinate-ascent structure: an E-step producing responsibilities, then a per-block conditional M-step), with two differences that only affect SPEED, never correctness:

  1. A component D1 reports as frozen (see FreezeRollupCache.is_frozen()) reuses last round’s cached per-datum log-density instead of recomputing it, and is excluded from the M-step (its model object carries forward unchanged) – so a round with only k of K components active costs O(k/K) of a full round’s seq_log_density work.

  2. Every round is objective-gated exactly like MonotonicEM: the candidate model’s objective is checked (again cache-aware, so this costs the same O(active fraction)) and the step is rejected – keeping the previous model – if the objective would decrease beyond accept_tolerance. This is what makes the returned history a real monotone-F receipt, not just an assumption.

Returns (final_model, history) where history[i] is round i’s FreezeRollupStats (including objective, the real Neal-Hinton F for that round, and n_log_density_evals, the wall-clock proxy this module optimizes).

Parameters:

  • enc_data (Any)

  • estimator (MixtureEstimator)

  • initial_model (MixtureDistribution)

  • max_its (int)

  • delta (float | None)

  • cache (FreezeRollupCache | None)

  • q_gain_tol (float)

  • weight_tol (float)

  • weight_delta_tol (float)

  • freeze_patience (int)

  • accept_tolerance (float)

Return type:

tuple[MixtureDistribution, list[FreezeRollupStats]]