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 an update_kind classification. D2 combines that structural signal with directly evaluated component scores: once a subtree is structurally frozen, or its real-data scores and collapsed weight have converged, 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: exact freeze/roll-up is a scheduling optimization only. It caches only structurally frozen components and never changes the EM target. Near-zero-weight heuristic freezing is a separate, explicit approximation: callers must opt in, its eligibility threshold is the declared approximation budget, and every heuristically frozen component is forcibly reactivated after a bounded number of rounds. Cache entries bind the complete encoded-data digest, component state digest, component index, and compute dtype, so reuse cannot cross datasets or scoring modes.

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 DensityMatrixProfile(evaluations, cache_hits, reused_columns, zero_weight_skips, evaluation_seconds, cache_hit_seconds, assembly_seconds, component_evaluation_seconds)[source]

Bases: object

Observed cache/evaluation work for one component-matrix construction.

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

Bases: object

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

n_log_density_evals is the model-work receipt 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 directly evaluated, transactionally monotone observed-data MLE or MAP objective for this round.

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)

  • n_approximate_frozen (int)

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, approximate_freezing=False, max_frozen_rounds=10)[source]

Bases: object

Per-mixture-component cache keyed by data, model state, index, and compute dtype.

Structurally frozen components are exact cache candidates. Heuristic freezing is disabled by default because skipping a non-neutral component’s M-step changes the optimization target. approximate_freezing=True opts into that behavior with weight_tol as the eligibility budget and max_frozen_rounds as a mandatory escape/recheck interval.

Parameters:

  • q_gain_tol (float)

  • weight_tol (float)

  • weight_delta_tol (float)

  • freeze_patience (int)

  • approximate_freezing (bool)

  • max_frozen_rounds (int)

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

bind_data(enc_data)[source]

Bind the cache and convergence bookkeeping to one encoded dataset digest.

Parameters:

enc_data (Any)

Return type:

bytes

record_component_scores(matrix)[source]

Record accepted real-data component scores and their maximum rowwise movement.

Parameters:

matrix (ndarray)

Return type:

None

is_frozen(idx, component, weight)[source]

Return the exact or explicitly approximate freeze verdict for idx.

A structurally "frozen" node is exact. Otherwise, only an opted-in approximate cache can freeze it, after both scores and a below-budget weight stabilize for freeze_patience rounds. Score/weight movement escapes immediately, and max_frozen_rounds forces periodic reactivation because plateau evidence is not proof that skipping the component’s M-step preserves the exact fixed point.

Parameters:
Return type:

bool

is_approximately_frozen(idx)[source]

Whether the current freeze verdict for idx used the opt-in heuristic.

Parameters:

idx (int)

Return type:

bool

component_log_density(idx, component, enc, *, frozen, compute_dtype=None, data_signature=None)[source]

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

A hit still verifies the component content signature, but avoids component.seq_log_density. Any model, data, index, or compute-dtype mismatch forces a recompute.

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]

is_compiled_em_eligible(model, estimator)[source]

Whether the local fused full-sweep adapter can execute this model/estimator pair.

Parameters:
Return type:

bool

compiled_em_step(enc_data, estimator, model, *, compute_dtype=None, cache=None)[source]

Run one full mixture E/M sweep through the reusable fused component path.

Parameters:

  • enc_data (Any)

  • estimator (MixtureEstimator)

  • model (MixtureDistribution)

  • compute_dtype (Any)

  • cache (FreezeRollupCache | None)

Return type:

tuple[MixtureDistribution, dict[str, Any]]

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, approximate_freezing=False, max_frozen_rounds=10, 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. By default, only structurally neutral components are frozen, so caching affects speed but not the EM target.

approximate_freezing=True additionally permits converged, near-zero-weight components to skip M-steps. That changes the reachable target and is therefore explicitly approximate. weight_tol is its eligibility budget; weight/score drift immediately escapes, and max_frozen_rounds forces a full component recheck even without detected drift. History reports how many components used this approximation each round.

Every round is objective-gated: a candidate that decreases the observed-data MLE/MAP objective beyond accept_tolerance is rejected.

Returns (final_model, history) where history[i] is round i’s FreezeRollupStats (including objective, the observed-data MLE/MAP objective for that round, and n_log_density_evals, the component-evaluation work receipt 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)

  • approximate_freezing (bool)

  • max_frozen_rounds (int)

  • accept_tolerance (float)

Return type:

tuple[MixtureDistribution, list[FreezeRollupStats]]