Fitting

EM Fitters

EM fitters for normix distributions.

Model knows math, fitter knows iteration.

BatchEMFitter is batch EM with a dual-loop architecture (lax.scan when both backends are JAX, otherwise a Python for-loop). IncrementalEMFitter is online / mini-batch EM with pluggable eta update rules.

class normix.fitting.em.EMResult(model, log_likelihoods, param_changes, n_iter, converged, elapsed_time, diverged=False, solver_failed=False)

Bases: object

Result of an EM fitting procedure.

converged, diverged, and solver_failed are separate. Tolerance met is converged=True. A non-finite normal-block change or objective is diverged=True: the last finite model is kept. A non-finite mixing distribution after the M-step is solver_failed=True, again with the last finite model kept. On the Python loop, SolverDidNotConverge or ValueError sets the same flag. Exhausting max_iter leaves all three false. They are not opposites.

Parameters:

• model (Any)

• log_likelihoods (Array | None)

• param_changes (Array)

• n_iter (int)

• converged (bool | None)

• elapsed_time (float)

• diverged (bool)

• solver_failed (bool)

model: Any

log_likelihoods: Array | None

param_changes: Array

n_iter: int

converged: bool | None

elapsed_time: float

diverged: bool = False

solver_failed: bool = False

class normix.fitting.em.BatchEMFitter(*, algorithm='em', max_iter=200, tol=0.001, verbose=0, regularization='none', e_step_backend='jax', m_step_backend='cpu', m_step_method='newton', eta_update=None, track_ll=False, m_step_kwargs=None)

Bases: object

Batch EM / MCECM algorithm with dual-loop architecture.

EM (default): E-step → M-step (all params) → regularize.

MCECM: E-step → M-step (normal params only) → regularize → E-step → M-step (subordinator only).

Convergence is the Aitken remaining gap of the mean log-likelihood \(\ell_n\) (nats per observation) from the E-step conjugacy identity; stop when \(\ell_\infty - \ell_{t+1} <\) tol. Three consecutive \(\ell\) values are required, so the earliest stop is n_iter=3. Hybrid-scale RMS on em_convergence_params() is recorded in EMResult.param_changes as a diagnostic, not the stop. Covariance regularisations leave the density unchanged, so they share this criterion. Under scalar-\(\tau\) MAP shrinkage the tracked sequence is the penalised objective.

Loop selection is automatic:

• lax.scan when both backends are ‘jax’, verbose <= 1, algorithm=’em’, and no eta_update rule

• Python for-loop otherwise

Parameters:

• algorithm (str) – ‘em’ (default) or ‘mcecm’.

• max_iter (int) – Maximum number of iterations.

• tol (float) – Convergence tolerance on the Aitken remaining gap of \(\ell_n\) in nats per observation. tol=1e-3 means the extrapolated remaining improvement is less than \(10^{-3}\) nats/obs.

• verbose (int) – 0 = silent, 1 = summary, 2 = per-iteration table.

• regularization (str) –

  Strategy applied after each M-step:

  • 'none' — no regularization.

  • 'det_sigma_one' — enforce \(|\Sigma| = 1\) (the original GH convention; equivalent to regularize_det_sigma(target_log_det=0)).

  • 'det_sigma_x' — enforce \(|\Sigma| = |\Sigma_0|\) where \(\Sigma_0\) is the dispersion of the initial model passed to fit(). Useful when comparing GH / FactorGH parameters against VG / NIG / NInvG, which leave \(|\Sigma|\) at the empirical scale.

  • 'a_eq_b' — rescale the GIG subordinator so \(a = b = \sqrt{ab}\). Trivial no-op for VG / NInvG / MultivariateNormal.

• e_step_backend (str) – ‘jax’ (default) or ‘cpu’.

• m_step_backend (str) – ‘jax’ or ‘cpu’ (default, faster for GIG).

• m_step_method (str) – ‘newton’ (default), ‘lbfgs’, or ‘bfgs’.

• eta_update (EtaUpdateRule or None) – Optional eta combination rule (e.g. Shrinkage(IdentityUpdate(), eta0, tau)). When set, the E-step output is transformed before the M-step.

• track_ll (bool) – When True, record the per-iteration conjugacy mean log-likelihood (the Aitken sequence) in EMResult.log_likelihoods without requiring verbose >= 1. These are \(\ell_n(\theta_t)\) at the E-step parameters, not a second Bessel pass through marginal_log_likelihood().

• m_step_kwargs (dict or None) – Extra keyword arguments forwarded verbatim to every m_step / m_step_subordinator call (in addition to backend and method). Used to thread estimand controls such as the VG alpha_min shape bound down to the subordinator’s from_expectation. Values must be static (e.g. Python floats) to stay compatible with the lax.scan path. None = no extras.

fit(model, X)

Run batch EM or MCECM. Auto-selects lax.scan or Python loop.

Parameters:

• model (NormalMixture) – Instance of a NormalMixture subclass, used as initial parameters.

• X (jax.Array) – Data array, shape (n, d).

Returns:

• EMResult with fitted model, convergence diagnostics, and timing.

• X is centered at its sample mean before any statistic is

• accumulated, and that mean is added back onto mu. gamma,

• Sigma, and the subordinator are unchanged by the shift. A

• shrinkage target equal to the initial model’s own η is rebuilt

• in these coordinates; any other target is translated as a

• population moment.

Return type:

EMResult

class normix.fitting.em.IncrementalEMFitter(*, eta_update=None, batch_size=256, max_steps=200, inner_iter=1, verbose=0, regularization='none', e_step_backend='jax', m_step_backend='cpu', m_step_method='newton')

Bases: object

Incremental EM with pluggable eta update rules.

Replaces OnlineEMFitter and MiniBatchEMFitter. Processes data in random mini-batches, applies an EtaUpdateRule to combine the running \(\eta\) with each batch estimate, then M-steps on the combined \(\eta\).

Parameters:

• eta_update (EtaUpdateRule) – How to combine running η with each batch estimate.

• batch_size (int) – Observations per batch.

• max_steps (int) – Number of batches to process (total budget).

• inner_iter (int) – 1 = online (default); >1 = fine-tuning on each batch.

• verbose (int) –

  0 = silent; 1 = periodic summary.

  Scan path (JAX backends only):

  verbose must be 0 so diagnostics do not rely on Python side effects each step.

• regularization (str) – Same options as BatchEMFitter — 'none' | 'det_sigma_one' | 'det_sigma_x' | 'a_eq_b'.

• e_step_backend (str) – Passed through to e_step / m_step.

• m_step_backend (str) – Passed through to e_step / m_step.

• m_step_method (str) – Passed through to e_step / m_step.

fit(model, X, *, key)

Run incremental EM. Returns EMResult.

X is centered at its sample mean for the whole stream, and that mean is added back onto mu. A shrinkage target equal to the initial model’s own η is rebuilt in these coordinates.

Parameters:

• X (Array)

• key (Array)

Return type:

EMResult

Solvers

Bregman divergence solvers.

Minimises f(θ) − θ·η over θ, where f is any convex function (e.g. the log-partition ψ for an exponential family). At the minimum ∇f(θ*) = η.

Public API

solve_bregman

Single starting point.

solve_bregman_multistart

Ordered starting points, stopping at the first finite converged JAX solve. CPU solves select the best result from a Python loop.

bregman_objective

Utility \(f(\theta) - \theta\cdot\eta\).

make_jit_newton_solver

Build a stable @jax.jit Newton solve specialised to a fixed (f, grad_fn, hess_fn, bounds). Repeated calls with matching shapes/dtypes hit the XLA cache, avoiding the per-call re-tracing that solve_bregman incurs from fresh closures.

Backends × methods

backend=’jax’, method=’newton’ custom lax.while_loop Newton, autodiff or analytical Hessian backend=’jax’, method=’lbfgs’ jaxopt LBFGSB (bounds native) or LBFGS (reparam) backend=’jax’, method=’bfgs’ jaxopt BFGS with reparameterization for bounds backend=’cpu’, method=’lbfgs’ scipy L-BFGS-B (native bounds) backend=’cpu’, method=’bfgs’ scipy BFGS; bounds via the same reparametrization backend=’cpu’, method=’newton’ scipy trust-exact; bounds via the same reparametrization

Gradient / Hessian sources

grad_fnθ → ∇f(θ). For backend=’cpu’: pure CPU (numpy) gradient.

For backend=’jax’, method=’newton’: JAX-traceable ∇ψ(θ). If None with backend=’cpu’: hybrid — jax.grad compiled → NumPy callbacks.

hess_fnθ → ∇²f(θ). Required for method=’newton’.

For backend=’cpu’: pure CPU (numpy) Hessian. For backend=’jax’, method=’newton’: JAX-traceable ∇²ψ(θ). If None with method=’newton’: jax.hessian of the full objective.

Both grad_fn and hess_fn operate in theta-space only. The solver handles all reparameterization internally via the chain rule.

class normix.fitting.solvers.BregmanResult(theta, fun, grad_norm, num_steps, converged, elapsed_time=0.0, kkt_residual=None, reason=None)

Bases: object

Result of a Bregman divergence minimization.

Scalar fields accept both Python and JAX types so the result can live inside a lax.scan carry without concretization errors.

Parameters:

• theta (Array)

• fun (Any)

• grad_norm (Any)

• num_steps (Any)

• converged (Any)

• elapsed_time (float)

• kkt_residual (Any)

• reason (Any)

theta: Array

fun: Any

grad_norm: Any

num_steps: Any

converged: Any

elapsed_time: float = 0.0

kkt_residual: Any = None

reason: Any = None

exception normix.fitting.solvers.SolverDidNotConverge

Bases: RuntimeError

Eager Bregman inversion did not meet the residual tolerance.

normix.fitting.solvers.bregman_objective(theta, eta, f)

f(θ) − θ·η — convex dual whose minimum gives ∇f(θ*) = η.

Parameters:

• theta (Array)

• eta (Array)

• f (Callable[[Array], Array])

Return type:

Array

normix.fitting.solvers.solve_bregman(f, eta, theta0, *, backend='jax', method='lbfgs', bounds=None, max_steps=500, tol=1e-10, grad_fn=None, hess_fn=None, verbose=0)

Minimise f(θ) − θ·η over θ.

Parameters:

• f (Callable) – Convex function θ → scalar (e.g. log-partition ψ).

• eta (jax.Array) – Target vector (e.g. expectation parameters η).

• theta0 (jax.Array) – Initial guess.

• backend ('jax' or 'cpu') – 'jax' is JIT-able; 'cpu' uses scipy and is not JIT-able.

• method ('lbfgs', 'bfgs', or 'newton')

• bounds (tuple of jax.Array, or None) – (lower, upper) pair, each shape (d,); None → unconstrained. JAX Newton, JAX BFGS, CPU BFGS, and CPU trust-exact enforce bounds by reparameterization. CPU L-BFGS-B takes them natively.

• max_steps (int) – Iteration budget.

• tol (float) – Convergence tolerance on ‖∇f(θ) − η‖∞. On a reparametrised CPU solve, SciPy gtol is the φ-gradient. converged is still the free θ residual.

• grad_fn (Callable, optional) – θ → ∇f(θ). For backend=’cpu’: must accept and return numpy arrays. For backend=’jax’, method=’newton’: must be JAX-traceable. If None with backend=’cpu’: falls back to jax.grad (hybrid mode).

• hess_fn (Callable, optional) – θ → ∇²f(θ). Required for method=’newton’. For backend=’cpu’: must accept numpy arrays and return numpy array. For backend=’jax’, method=’newton’: must be JAX-traceable. If None with method=’newton’: jax.hessian of the full objective is used.

• verbose (int) – 0 = silent, >= 1 = print summary after solve. Summaries are skipped while tracing.

Return type:

BregmanResult

normix.fitting.solvers.solve_bregman_multistart(f, eta, theta0_batch, *, backend='jax', method='lbfgs', bounds=None, max_steps=500, tol=1e-10, grad_fn=None, hess_fn=None, verbose=0)

Solve from ordered starting points, preferring finite convergence.

Parameters:

• theta0_batch (jax.Array or list of jax.Array) – Nonempty (K, dim) array or list of arrays. JAX tries starts in order and stops at the first finite converged result. If none converges, it returns the lowest finite objective. Put preferred starts first. CPU solves use a Python loop over all starts.

• verbose (int) – 0 = silent, >= 1 = print summary. Summaries are skipped while tracing.

• f (Callable[[Array], Array])

• eta (Array)

• backend (str)

• method (str)

• bounds (Tuple[Array, Array] | None)

• max_steps (int)

• tol (float)

• grad_fn (Callable | None)

• hess_fn (Callable | None)

Return type:

BregmanResult

normix.fitting.solvers.make_jit_newton_solver(f, grad_fn, hess_fn, bounds=None)

Build a @jax.jit-decorated Newton solver specialised to one problem.

The returned callable has signature solve(eta, theta0, max_steps=20, tol=1e-10) -> (theta_opt, fun, grad_norm, converged) where max_steps and tol are static arguments.

All distribution-level inputs (f, grad_fn, hess_fn, bounds) are baked into the closure at construction time. Repeated calls with the same array shapes and dtypes therefore reuse the compiled XLA executable.

Use this in EM hot paths where solve_bregman would otherwise build a fresh Python closure on every call and force JAX to re-trace the same Newton kernel on each iteration.

Parameters:

• f (Callable) – Convex objective ψ(θ) → scalar (must be JAX-traceable).

• grad_fn (Callable) – \(\nabla\psi(\theta) \to \mathbb{R}^d\).

• hess_fn (Callable) – \(\nabla^2\psi(\theta) \to \mathbb{R}^{d\times d}\).

• bounds (tuple or None) – (lower, upper), each shape \((d,)\), or None.

Returns:

Jit-compiled Newton solver. Returns a 4-tuple of JAX arrays (theta, fun, grad_norm, converged); wrap in BregmanResult externally if needed.

Return type:

Callable

Eta Parametrization

NormalMixtureEta — expectation parametrization for normal variance-mean mixtures.

The six fields are the batch-averaged sufficient statistics \(\hat\eta = \frac{1}{n}\sum_i E[t(X_i, Y_i) \mid X_i]\), in the theory order used in Shrinkage with Penalized Likelihood and Factor Analysis for Generalized Hyperbolic Distributions:

\[s_1 = E[Y^{-1}], \;\; s_2 = E[Y], \;\; s_3 = E[\log Y], \;\; s_4 = E[X], \;\; s_5 = E[X / Y], \;\; s_6 = E[X X^\top / Y].\]

This is the expectation parametrization of JointNormalMixture. GIG arrays use the last-axis order (E_log_Y, E_inv_Y, E_Y). The private packing helpers convert between these layouts by name.

class normix.fitting.eta.EStepResult(eta, log_lik)

Bases: Module

E-step output: aggregated \(\eta\) and mean log-likelihood.

\(\ell_n\) is not a field of NormalMixtureEta — affine_combine / shrinkage would mix it. It is computed from the E-step’s posterior GIG log-partition (already in BesselMoments.log_k) via conjugacy, with no extra Bessel calls and no cache on the model.

Parameters:

• eta (Any)

• log_lik (Array)

eta: Any

log_lik: Array

mean log-likelihood \(\ell_n(\theta)\) at the E-step parameters, nats per observation

class normix.fitting.eta.NormalMixtureEta(E_inv_Y, E_Y, E_log_Y, E_X, E_X_inv_Y, E_XXT_inv_Y)

Bases: Module

Aggregated expectation parameters for normal variance-mean mixtures.

Fields are stored in theory order (s_1, …, s_6): the first six statistics are shared with FactorMixtureStats so that shrinkage targets, weights, and tests written for the standard family transfer unchanged.

Parameters:

• E_inv_Y (Array)

• E_Y (Array)

• E_log_Y (Array)

• E_X (Array)

• E_X_inv_Y (Array)

• E_XXT_inv_Y (Array)

E_inv_Y: Array

scalar; \(s_1 = \frac{1}{n}\sum_i E[1/Y_i \mid X_i]\)

E_Y: Array

scalar; \(s_2 = \frac{1}{n}\sum_i E[Y_i \mid X_i]\)

E_log_Y: Array

scalar; \(s_3 = \frac{1}{n}\sum_i E[\log Y_i \mid X_i]\)

E_X: Array

shape \((d,)\); \(s_4 = \frac{1}{n}\sum_i X_i\)

E_X_inv_Y: Array

shape \((d,)\); \(s_5 = \frac{1}{n}\sum_i X_i \, E[1/Y_i \mid X_i]\)

E_XXT_inv_Y: Array

shape \((d, d)\); \(s_6 = \frac{1}{n}\sum_i X_i X_i^\top E[1/Y_i \mid X_i]\)

class normix.fitting.eta.FactorMixtureStats(E_inv_Y, E_Y, E_log_Y, E_X, E_X_inv_Y, E_XXT_inv_Y, E_XZT_inv_sqrtY, E_Z_inv_sqrtY, E_Z_sqrtY, E_ZZT)

Bases: Module

Aggregated expectation parameters for factor-analysis mixtures.

Fields are stored in theory order (s_1, …, s_{10}) from Factor Analysis for Generalized Hyperbolic Distributions. The first six are identical to NormalMixtureEta (so shrinkage targets, η-update rules, and weight pytrees designed for the standard family broadcast onto the factor family without modification). The four extra fields involve the latent factor \(Z\).

Parameters:

• E_inv_Y (Array)

• E_Y (Array)

• E_log_Y (Array)

• E_X (Array)

• E_X_inv_Y (Array)

• E_XXT_inv_Y (Array)

• E_XZT_inv_sqrtY (Array)

• E_Z_inv_sqrtY (Array)

• E_Z_sqrtY (Array)

• E_ZZT (Array)

E_inv_Y: Array

scalar; \(s_1 = \frac{1}{n}\sum_i E[1/Y_i \mid X_i]\)

E_Y: Array

scalar; \(s_2 = \frac{1}{n}\sum_i E[Y_i \mid X_i]\)

E_log_Y: Array

scalar; \(s_3 = \frac{1}{n}\sum_i E[\log Y_i \mid X_i]\)

E_X: Array

shape \((d,)\); \(s_4 = \frac{1}{n}\sum_i X_i\)

E_X_inv_Y: Array

shape \((d,)\); \(s_5 = \frac{1}{n}\sum_i X_i \, E[1/Y_i \mid X_i]\)

E_XXT_inv_Y: Array

shape \((d, d)\); \(s_6 = \frac{1}{n}\sum_i X_i X_i^\top E[1/Y_i \mid X_i]\)

E_XZT_inv_sqrtY: Array

shape \((d, r)\); \(s_7 = \frac{1}{n}\sum_i E[X_i Z_i^\top Y_i^{-1/2} \mid X_i]\)

E_Z_inv_sqrtY: Array

shape \((r,)\); \(s_8 = \frac{1}{n}\sum_i E[Z_i Y_i^{-1/2} \mid X_i]\)

E_Z_sqrtY: Array

shape \((r,)\); \(s_9 = \frac{1}{n}\sum_i E[Z_i Y_i^{1/2} \mid X_i]\)

E_ZZT: Array

shape \((r, r)\); \(s_{10} = \frac{1}{n}\sum_i E[Z_i Z_i^\top \mid X_i]\)

normix.fitting.eta.affine_combine(eta_prev, eta_new, b, c, a=None)

Affine combination \(\eta_t = a + b\,\eta_{t-1} + c\,\hat\eta\).

The weights b and c may be:

• scalar (Python number or 0-d jax.Array) — broadcast to every leaf of eta;

• stats-shape pytree (same type as eta_prev / eta_new) — block-diagonal weighting; leaf-wise multiply;

• callable η → η — arbitrary linear operator on η (e.g. an eqx.nn.Linear wrapped to operate on a flattened pytree).

The shift a is either None (zero) or a stats-shape pytree.

Parameters:

• eta_prev – Running state \(\eta_{t-1}\).

• eta_new – New batch estimate \(\hat\eta\).

• b (float | Array | NormalMixtureEta | Callable[[...], NormalMixtureEta]) – Weight on previous state.

• c (float | Array | NormalMixtureEta | Callable[[...], NormalMixtureEta]) – Weight on new estimate.

• a – Additive shift (e.g. shrinkage prior). None means zero.

Eta Update Rules

Eta update rules for incremental and penalised EM.

Two-layer abstraction

The most general rule is

\[\eta_t = \mathrm{rule}(\eta_{t-1},\, \hat\eta_{\text{batch}}),\]

implemented by EtaUpdateRule via __call__. This leaves room for non-affine predictors (e.g. an MLP that maps (η_{t-1}, η̂) → η_t) without another API revision.

Most rules in this module are affine:

\[\eta_t = a + b\,\eta_{t-1} + c\,\hat\eta_{\text{batch}},\]

implemented by AffineRule, which delegates to AffineRule.weights() returning (a, b, c, state) and combines via affine_combine().

All rules are equinox.Module pytrees so their hyperparameters (e.g. tau0, w, tau) are JAX array leaves — JIT-compatible and differentiable for future meta-learning of step-size schedules.

class normix.fitting.eta_rules.EtaUpdateRule

Bases: Module

Abstract base for eta update rules.

The fitter only knows __call__(); whether a concrete rule is affine, a combinator, or an ML-style predictor is invisible at the call site.

Subclasses override __call__(). Rules with no per-step memory inherit initial_state() returning an empty dict.

initial_state()

Return type:

Dict

class normix.fitting.eta_rules.AffineRule

Bases: EtaUpdateRule

Specialisation: \(\eta_t = a + b\,\eta_{t-1} + c\,\hat\eta\).

Subclasses implement weights() instead of __call__(). The base class provides a single __call__() that delegates to weights() and runs the combination through affine_combine().

abstractmethod weights(step, batch_size, state)

Return (a, b, c, updated_state).

Parameters:

state (Dict)

Return type:

Tuple[NormalMixtureEta | None, Array, Array, Dict]

class normix.fitting.eta_rules.IdentityUpdate

Bases: AffineRule

Pass-through: \(\eta_t = \hat\eta\) (standard batch EM).

weights(step, batch_size, state)

Return (a, b, c, updated_state).

class normix.fitting.eta_rules.RobbinsMonroUpdate(tau0=10.0)

Bases: AffineRule

Robbins–Monro: \(c = 1/(\tau_0 + t)\), \(b = 1 - c\).

Parameters:

tau0 (float) – Initial step-size denominator (higher → slower adaptation).

tau0: Array

weights(step, batch_size, state)

Return (a, b, c, updated_state).

class normix.fitting.eta_rules.SampleWeightedUpdate

Bases: AffineRule

Incremental mean: \(b = n/(n+m)\), \(c = m/(n+m)\).

Tracks cumulative sample count n; each batch contributes m.

weights(step, batch_size, state)

Return (a, b, c, updated_state).

initial_state()

class normix.fitting.eta_rules.EWMAUpdate(w=0.1)

Bases: AffineRule

Exponentially weighted moving average: \(b = 1-w\), \(c = w\).

Parameters:

w (float) – Weight on the new batch (0 < w ≤ 1).

w: Array

weights(step, batch_size, state)

Return (a, b, c, updated_state).

class normix.fitting.eta_rules.Shrinkage(base, eta0, tau=0.5)

Bases: EtaUpdateRule

Shrinkage combinator: pull the running η toward a prior on top of any base rule.

\[\eta_t = \frac{\tau}{1+\tau} \odot \eta_0 + \frac{1}{1+\tau} \odot \mathrm{base}(\eta_{t-1}, \hat\eta),\]

where \(\odot\) is per-field multiplication (broadcast for scalar \(\tau\)).

The combinator composes with any base rule — affine (IdentityUpdate, RobbinsMonroUpdate, EWMAUpdate, SampleWeightedUpdate) or non-affine (e.g. an MLP-based predictor) — without losing the base rule’s state.

Parameters:

• base (EtaUpdateRule) – Base rule that produces the unshrunk update \(\mathrm{base}(\eta_{t-1}, \hat\eta)\).

• eta0 (NormalMixtureEta) – Prior expectation parameters (shrinkage target). Must be a complete stats pytree of the same type as eta_prev / eta_new.

• tau (float, jax.Array, or NormalMixtureEta) –

  Shrinkage strength.

  • scalar — uniform shrinkage on every sufficient statistic; matches the penalised-MLE in Shrinkage with Penalized Likelihood.

  • stats pytree (e.g. NormalMixtureEta with scalar leaves) — per-field shrinkage. Setting all but one leaf to 0 shrinks only that statistic (e.g. Σ alone via the E_XXT_inv_Y field).

Examples

Batch EM with uniform shrinkage:

rule = Shrinkage(IdentityUpdate(), eta0, tau=0.5)

Batch EM with Σ-only shrinkage:

tau_pytree = NormalMixtureEta(
    E_inv_Y=0.0, E_Y=0.0, E_log_Y=0.0,
    E_X=jnp.zeros(d), E_X_inv_Y=jnp.zeros(d),
    E_XXT_inv_Y=jnp.full((d, d), 0.5),
)
rule = Shrinkage(IdentityUpdate(), eta0, tau=tau_pytree)

Robbins–Monro online + shrinkage:

rule = Shrinkage(RobbinsMonroUpdate(tau0=10.0), eta0, tau=0.1)

Notes

The state pytree is owned by base: initial_state() and __call__() thread the base rule’s state unchanged through the shrinkage step.

See also

normix.fitting.shrinkage_targets

helpers for building eta0.

base: EtaUpdateRule

eta0: Any

tau: Any

initial_state()

Return type:

Dict

class normix.fitting.eta_rules.AffineUpdate(a=None, b=0.0, c=1.0)

Bases: AffineRule

User-defined constant \((a, b, c)\).

All three coefficients are pytree values — b and c are scalar jax.Array leaves, a is an optional NormalMixtureEta. For time-varying schedules, subclass EtaUpdateRule directly.

Parameters:

• a (NormalMixtureEta or None) – Additive shift (e.g. scaled prior).

• b (float) – Weight on previous state.

• c (float) – Weight on new batch.

a: NormalMixtureEta | None

b: Array

c: Array

weights(step, batch_size, state)

Return (a, b, c, updated_state).

Shrinkage Targets

Shrinkage target builders for penalised EM.

The Shrinkage combinator pulls the running expectation parameters toward a prior \(\eta_0\). The helpers below construct that prior from a fitted (or moment-initialised) NormalMixture, optionally substituting a custom dispersion \(\Sigma_0\).

Per Shrinkage with Penalized Likelihood (Shrunk Sufficient Statistics), the prior expectation parameters are

\[\begin{split}s_1 &= E[Y^{-1}\mid\theta_0],\quad s_2 = E[Y\mid\theta_0],\quad s_3 = E[\log Y\mid\theta_0] \\ s_4 &= \mu_0 + \gamma_0\,E[Y\mid\theta_0] \\ s_5 &= \mu_0\,E[Y^{-1}\mid\theta_0] + \gamma_0 \\ s_6 &= \Sigma_0 + \mu_0\mu_0^\top E[Y^{-1}\mid\theta_0] + \gamma_0\gamma_0^\top E[Y\mid\theta_0] + \mu_0\gamma_0^\top + \gamma_0\mu_0^\top.\end{split}\]

All four constructors return a complete six-field NormalMixtureEta. The Σ-only variants reuse the model’s own \((\mu, \gamma, p, a, b)\) to fill the other five fields, keeping the public contract simple (”eta0 is always a full prior”) while the user’s per-field tau selects which fields are actually shrunk.

normix.fitting.shrinkage_targets.eta0_from_model(model)

Prior \(\eta_0\) equal to the model’s current expectation parameters.

Equivalent to model.compute_eta_from_model().

Parameters:

model (NormalMixture) – Source model; the prior reuses its \((\mu, \gamma, \Sigma)\) and subordinator parameters.

Returns:

Six-field expectation pytree.

Return type:

NormalMixtureEta

normix.fitting.shrinkage_targets.eta0_isotropic(model, sigma2)

Prior with an isotropic dispersion \(\Sigma_0 = \sigma^2 I_d\).

Parameters:

• model (NormalMixture) – Source model; provides \((\mu, \gamma, p, a, b)\).

• sigma2 (float) – Common variance for the isotropic prior. Must be positive.

Returns:

Six-field expectation pytree.

Return type:

NormalMixtureEta

normix.fitting.shrinkage_targets.eta0_diagonal(model, diag)

Prior with a diagonal dispersion \(\Sigma_0 = \mathrm{diag}(\text{diag})\).

Parameters:

• model (NormalMixture) – Source model; provides \((\mu, \gamma, p, a, b)\).

• diag (jax.Array) – Diagonal entries of \(\Sigma_0\), shape \((d,)\). Must be positive.

Returns:

Six-field expectation pytree.

Return type:

NormalMixtureEta

normix.fitting.shrinkage_targets.eta0_with_sigma(model, Sigma0)

Prior \(\eta_0\) reusing model parameters with a custom \(\Sigma_0\).

Substitutes Sigma0 for the model’s covariance in the \(s_6\) term while keeping \(\mu, \gamma\) and the subordinator expectations from model. This is the building block for “shrink Σ only” workflows: combine with a per-field tau that is non-zero only on E_XXT_inv_Y.

Parameters:

• model (NormalMixture) – Source model; provides \((\mu, \gamma, p, a, b)\).

• Sigma0 (jax.Array) – Prior dispersion to embed in \(s_6\), shape \((d, d)\). Must be positive semi-definite (not checked).

Returns:

Six-field expectation pytree with E_XXT_inv_Y rebuilt from Sigma0.

Return type:

NormalMixtureEta