Fitting

Contents

Fitting#

EM Fitters#

EM fitters for normix distributions.

Model knows math, fitter knows iteration.

BatchEMFitter — standard batch EM with dual-loop architecture:

lax.scan (JIT-able) or Python for-loop (CPU-compatible)

IncrementalEMFitter — 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)[source]#

Bases: object

Result of an EM fitting procedure.

converged and diverged are separate because batch EM has three outcomes: tolerance met (converged=True), a non-finite iterate was reverted (diverged=True), or max_iter was exhausted with neither (both False). They are not opposites — converged=False covers both divergence and a finite but not-yet-converged stop.

Parameters:
model: Any#
log_likelihoods: Array | None#
param_changes: Array#
n_iter: int#
converged: bool | None#
elapsed_time: float#
diverged: 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)[source]#

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 measured by hybrid-scale RMS parameter change in the normal parameters (mu, gamma, L_Sigma), excluding subordinator (GIG) parameters. Per leaf the change is rms(new - old) / (1 + rms(old)) with rms(v) = ||v||_2 / sqrt(m) and m = v.size. Likelihood is not used for stopping (optional LL traces remain diagnostics only).

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 max hybrid-scale RMS parameter change (rms(Δ) / (1 + rms(θ))). Roughly dimension-free: tol=1e-3 means a typical coordinate moved by ~0.1% of its natural scale.

  • 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 per-iteration marginal log-likelihood in EMResult.log_likelihoods without requiring verbose >= 1.

  • 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)[source]#

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

Parameters:

  • model (NormalMixture subclass (used as initial parameters))

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

Return type:

EMResult with fitted model, convergence diagnostics, and timing.

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')[source]#

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)[source]#

Run incremental EM. Returns EMResult.

Parameters:
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 multiple starting points (vmap for JAX Newton;

for-loop for quasi-Newton and CPU)

bregman_objective utility: f(θ) − θ·η 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.scan 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 backend=’cpu’, method=’bfgs’ scipy BFGS backend=’cpu’, method=’newton’ scipy trust-exact with Hessian

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)[source]#

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: int#
converged: Any#
elapsed_time: float = 0.0#
normix.fitting.solvers.bregman_objective(theta, eta, f)[source]#

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

Parameters:
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)[source]#

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

Parameters:

  • f (convex function θ → scalar (e.g. log-partition ψ))

  • eta (target vector (e.g. expectation parameters η))

  • theta0 (initial guess)

  • backend ('jax' (JIT-able) or 'cpu' (scipy, not JIT-able))

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

  • bounds (tuple of jax.Array, or None) – (lower, upper) pair, each shape (d,); None → unconstrained. For backend=’jax’: enforced via reparameterization. For backend=’cpu’: converted to scipy format internally.

  • max_steps (iteration budget)

  • tol (convergence tolerance on ‖∇f(θ) − η‖∞)

  • grad_fn (θ → ∇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 (θ → ∇²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.

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)[source]#

Run solve_bregman from multiple starting points; return the best result.

Parameters:

  • theta0_batch ((K, dim) jax.Array for backend='jax', method='newton') – (parallel via vmap); list of arrays otherwise (sequential for-loop).

  • verbose (int) – 0 = silent, >= 1 = print summary.

  • 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)[source]#

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 is a static argument (required by lax.scan).

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 (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.

class normix.fitting.eta.NormalMixtureEta(E_inv_Y, E_Y, E_log_Y, E_X, E_X_inv_Y, E_XXT_inv_Y)[source]#

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#

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)[source]#

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#

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)[source]#

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 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[source]#

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()[source]#
Return type:

Dict

class normix.fitting.eta_rules.AffineRule[source]#

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)[source]#

Return (a, b, c, updated_state).

Parameters:

state (Dict)

Return type:

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

class normix.fitting.eta_rules.IdentityUpdate[source]#

Bases: AffineRule

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

weights(step, batch_size, state)[source]#

Return (a, b, c, updated_state).

class normix.fitting.eta_rules.RobbinsMonroUpdate(tau0=10.0)[source]#

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)[source]#

Return (a, b, c, updated_state).

class normix.fitting.eta_rules.SampleWeightedUpdate[source]#

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)[source]#

Return (a, b, c, updated_state).

initial_state()[source]#
class normix.fitting.eta_rules.EWMAUpdate(w=0.1)[source]#

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)[source]#

Return (a, b, c, updated_state).

class normix.fitting.eta_rules.Shrinkage(base, eta0, tau=0.5)[source]#

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 stats pytree) –

    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()[source]#
Return type:

Dict

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

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)[source]#

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)[source]#

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)[source]#

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)[source]#

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)[source]#

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

Contents