Mixtures

Contents

Mixtures#

Base classes shared by every mixture distribution in Distributions.

JointNormalMixture#

JointNormalMixture — abstract exponential family for normal variance-mean mixtures.

Joint distribution \(f(x, y)\):

\[X \mid Y \sim \mathcal{N}(\mu + \gamma y,\; \Sigma y), \quad Y \sim \text{subordinator (GIG, Gamma, InvGamma, InvGaussian)}\]

Sufficient statistics:

\[t(x, y) = [\log y,\; 1/y,\; y,\; x,\; x/y,\; \mathrm{vec}(xx^\top/y)]\]

Natural parameters:

\[\theta_1 = p_{\mathrm{sub}} - 1 - d/2, \quad \theta_2 = -(b_{\mathrm{sub}}/2 + \tfrac{1}{2}\mu^\top\Sigma^{-1}\mu) < 0\]
\[\theta_3 = -(a_{\mathrm{sub}}/2 + \tfrac{1}{2}\gamma^\top\Sigma^{-1}\gamma) < 0, \quad \theta_4 = \Sigma^{-1}\gamma, \quad \theta_5 = \Sigma^{-1}\mu, \quad \theta_6 = -\tfrac{1}{2}\mathrm{vec}(\Sigma^{-1})\]

For a GIG subordinator, \(\theta_2,\theta_3\) combine \(-b/2,-a/2\) with the normal quadratic forms so that \(\theta^{\top} t\) matches \(-(a y + b/y)/2\) from \(f_Y\) plus the \(y\)-dependent terms from \(f_{X\mid Y}\).

Log-partition:

\[\psi = \psi_{\mathrm{sub}}(p, a, b) + \tfrac{1}{2}\log|\Sigma| + \mu^\top\Sigma^{-1}\gamma\]

Expectation parameters (EM E-step quantities):

\[\eta_1 = E[\log Y], \quad \eta_2 = E[1/Y], \quad \eta_3 = E[Y]\]
\[\eta_4 = E[X] = \mu + \gamma E[Y], \quad \eta_5 = E[X/Y] = \mu E[1/Y] + \gamma\]
\[\eta_6 = E[XX^\top/Y] = \Sigma + \mu\mu^\top E[1/Y] + \gamma\gamma^\top E[Y] + \mu\gamma^\top + \gamma\mu^\top\]

EM M-step closed-form (let \(D = 1 - E[1/Y] \cdot E[Y]\)):

\[\mu = \frac{E[X] - E[Y] E[X/Y]}{D}, \quad \gamma = \frac{E[X/Y] - E[1/Y] E[X]}{D}\]
\[\Sigma = E[XX^\top/Y] - E[X/Y]\mu^\top - \mu E[X/Y]^\top + E[1/Y]\mu\mu^\top - E[Y]\gamma\gamma^\top\]
class normix.mixtures.joint.JointNormalMixture(mu, gamma, L_Sigma)[source]#

Bases: ExponentialFamily

Abstract joint distribution \(f(x, y)\) for normal variance-mean mixtures.

Stored: mu (d,), gamma (d,), L_Sigma (d×d lower Cholesky of \(\Sigma\)). Subordinator parameters defined by concrete subclasses.

Parameters:
mu: Array#
gamma: Array#
L_Sigma: Array#
abstractmethod subordinator()[source]#

Return the fitted subordinator distribution.

Return type:

ExponentialFamily

to_joint_generalized_hyperbolic(*, boundary_eps=0.0)[source]#

Exact embedding into JointGeneralizedHyperbolic.

Lifts the subordinator via to_gig() and keeps the Normal block (\(\mu, \gamma, L_\Sigma\)) unchanged. boundary_eps = 0 stores the Gamma/InvGamma boundary exactly (required for divergence gauges).

Parameters:

boundary_eps (float)

property d: int#
sigma()[source]#

Covariance matrix \(\Sigma = L_\Sigma L_\Sigma^\top\).

Return type:

Array

log_det_sigma()[source]#

\(\log|\Sigma| = 2\sum_i \log L_{ii}\), via Cholesky diagonal.

Return type:

Array

rvs(n, seed=42)[source]#

Sample \((X, Y)\) from the joint distribution via JAX PRNG.

Returns:

  • X (jax.Array) – Shape (n, d).

  • Y (jax.Array) – Shape (n,).

Parameters:
Return type:

Tuple[Array, Array]

log_prob_joint(x, y)[source]#

\(\log f(x, y) = \log f(x\mid y) + \log f_Y(y)\).

\[\log f(x\mid y) = -\tfrac{d}{2}\log(2\pi) - \tfrac{1}{2}\log|\Sigma| - \tfrac{d}{2}\log y - \frac{1}{2y}\|L^{-1}(x-\mu)\|^2 + \gamma^\top\Sigma^{-1}(x-\mu) - \tfrac{y}{2}\gamma^\top\Sigma^{-1}\gamma\]

\(\log f_Y(y)\) from subordinator.

Parameters:
Return type:

Array

conditional_expectations(x)[source]#

Compute \(E[g(Y)\mid X=x]\) for the EM E-step.

The posterior \(Y\mid X=x\) is

\[\mathrm{GIG}\!\left(p - \tfrac{d}{2},\; a + \gamma^\top\Sigma^{-1}\gamma,\; b + (x-\mu)^\top\Sigma^{-1}(x-\mu)\right),\]

with the family-specific prior \((p, a, b)\) resolved by _posterior_gig_params(). Returns a dict with keys E_log_Y, E_inv_Y, E_Y, plus psi_post (posterior GIG log-partition) and zw (\(\gamma^\top\Sigma^{-1}(x-\mu)\)) for the conjugacy identity \(\log f(x) = \log h(x) + \psi_{\mathrm{post}} - \psi\).

Parameters:

x (Array)

Return type:

Dict[str, Array]

static sufficient_statistics(xy)[source]#

\(t(x,y) = [\log y,\; 1/y,\; y,\; x,\; x/y,\; \mathrm{vec}(xx^\top/y)]\).

Input: flat vector \([x_1,\ldots,x_d, y]\).

Parameters:

xy (Array)

Return type:

Array

static log_base_measure(xy)[source]#

\(\log h(x)\) for a single unbatched observation.

Parameters:

xy (Array)

Return type:

Array

classmethod from_expectation(eta, **kwargs)[source]#

Construct from expectation parameters \(\eta\).

Two input forms are supported:

  • NormalMixtureEta — the natural η pytree of the joint normal-variance-mean mixture; uses the closed-form M-step (\(\mu, \gamma, \Sigma\) analytical; subordinator via _subordinator_from_eta()).

  • flat jax.Array — the generic Bregman solver inherited from ExponentialFamily.

The pytree path is the canonical η→θ map for these distributions: it is exact (no Bregman iterations on the normal block) and uses the subordinator’s own from_expectation (closed-form for Gamma / InverseGamma / InverseGaussian; numerical for GIG).

Parameters:
  • eta (NormalMixtureEta or jax.Array) – Expectation parameters.

  • **kwargs – For the pytree path: forwarded to _subordinator_from_eta() (e.g. backend, method, maxiter, theta0 for warm-starting GIG). For the flat-array path: forwarded to the parent solver.

Return type:

JointNormalMixture

NormalMixture#

Marginal mixture base classes.

MarginalMixture is the abstract interface fitters and the divergences module depend on. NormalMixture is the full-covariance implementation, owning a JointNormalMixture. A factor-analysis implementation lives in the sibling class FactorNormalMixture (see Mixture Architecture).

NormalMixture provides:

  • log_prob(x) — closed-form marginal log-density

  • e_step(X) — jax.vmap() over conditional expectations

  • m_step(X, expectations) — returns new NormalMixture

  • fit(X, ...) — convenience EM fitting with multi-start

class normix.mixtures.marginal.MarginalMixture[source]#

Bases: Module

Abstract interface for marginal mixtures used by the EM fitters.

Concrete subclasses pick the storage of the Gaussian dispersion (full Cholesky in NormalMixture, low-rank-plus-diagonal in FactorNormalMixture) and the type of the EM expectation pytree (NormalMixtureEta vs. FactorMixtureStats).

The fitter depends only on this contract; it does not know which storage form a model uses.

abstractmethod log_prob(x)[source]#

Marginal \(\log f(x)\) for a single observation.

Parameters:

x (Array)

Return type:

Array

pdf(x)[source]#

Marginal \(f(x)\) for a single observation.

Parameters:

x (Array)

Return type:

Array

marginal_log_likelihood(X)[source]#

Mean log-likelihood over a dataset.

Parameters:

X (Array)

Return type:

Array

abstractmethod e_step(X, *, backend='jax')[source]#

E-step: EStepResult (eta, log_lik).

Parameters:
Return type:

Any

abstractmethod m_step(eta, **kwargs)[source]#

Full M-step: updates all parameters; returns a new model.

Parameters:

eta (Any)

Return type:

MarginalMixture

abstractmethod m_step_normal(eta)[source]#

M-step for normal parameters only (MCECM cycle 1).

Parameters:

eta (Any)

Return type:

MarginalMixture

abstractmethod m_step_subordinator(eta, **kwargs)[source]#

M-step for subordinator parameters only (MCECM cycle 2).

Parameters:

eta (Any)

Return type:

MarginalMixture

abstractmethod compute_eta_from_model()[source]#

Reconstruct the expectation pytree from the model’s own parameters.

Return type:

Any

abstractmethod em_convergence_params()[source]#

Pytree whose leaf-wise hybrid-RMS change is the EM diagnostic.

Stopping uses the Aitken remaining gap of \(\ell_n\); this pytree feeds param_changes only. Subordinator parameters are excluded. For full-covariance models this is (mu, gamma, L_Sigma); for factor-analysis models it is (mu, gamma, F F^T + D) to sidestep the rotational gauge of \(F\).

Return type:

Any

abstractmethod subordinator()[source]#

Return the subordinator distribution \(Y\).

fit(X, *, alpha_min=None, **fitter_kwargs)[source]#

Fit using self as initialisation. Returns EMResult.

Keyword arguments other than alpha_min are forwarded to BatchEMFitter (algorithm, track_ll, eta_update, m_step_kwargs, backends, …). Family-specific defaults come from _fit_defaults() and are overridden by any explicitly supplied kwarg.

Parameters:
  • alpha_min (float or {'density', 'inverse_moment'}, optional) – Opt-in lower bound on the Gamma subordinator shape \(\alpha\) (Variance Gamma only — the unique family whose marginal likelihood is unbounded at \(x=\mu\)). Restricts the estimator to a region where that degeneracy cannot occur; the ghyp “fix-\(\lambda\)” analogue. None (default) leaves \(\alpha\) unconstrained. A float is used directly; the \(d\)-aware sentinels resolve to \(d/2 + \varepsilon\) ('density' — marginal density bounded) or \(d/2 + 1 + \varepsilon\) ('inverse_moment' — \(E[1/Y\mid x]\) also bounded), with \(\varepsilon=\) ALPHA_MIN_MARGIN. Has no effect on NInvG / NIG / GH (their prior \(b>0\) keeps the likelihood bounded for every \(\alpha\)). When set, it is merged into m_step_kwargs and wins over any alpha_min already present there.

  • X (jax.Array)

Return type:

EMResult

class normix.mixtures.marginal.NormalMixture(joint)[source]#

Bases: MarginalMixture

Marginal \(f(x) = \int_0^\infty f(x,y)\,dy\) for a normal variance-mean mixture.

Not an exponential family. Owns a JointNormalMixture (which is). The classical parameters \((\mu, \gamma, \Sigma, \text{subordinator})\) are forwarded as read-only properties; use replace() to obtain a new model with updated parameters (modules are immutable).

property joint: JointNormalMixture#
property d: int#
property mu: Array#

\(\mu\) — location parameter (forwarded from the joint).

property gamma: Array#

\(\gamma\) — skewness parameter (forwarded from the joint).

property L_Sigma: Array#

Lower Cholesky factor of \(\Sigma\) (forwarded from the joint).

sigma()[source]#

Dispersion \(\Sigma = L_\Sigma L_\Sigma^\top\) (forwarded from the joint).

Distinct from cov(), which returns the marginal covariance \(E[Y]\,\Sigma + \mathrm{Var}[Y]\,\gamma\gamma^\top\).

Return type:

Array

subordinator()[source]#

Subordinator distribution \(Y\) (forwarded from the joint).

log_det_sigma()[source]#

\(\log|\Sigma|\) (forwarded from the joint).

Return type:

Array

project(w)[source]#

Return the univariate normal mixture \(w^\top X\) as a Univariate* instance.

Parameters:

w (Array)

Return type:

_UnivariateNormalMixtureMixin

mean()[source]#

\(E[X] = \mu + \gamma E[Y]\).

A zero coordinate of \(\gamma\) does not form \(0\cdot\infty\). That coordinate equals \(\mu_i\) iff \(E[\sqrt{Y}]<\infty\) (NInvG / Student-\(t\) with \(\nu=2\alpha\): \(\alpha>1/2\)); otherwise \(+\infty\). Cauchy (\(\alpha=1/2\), \(\gamma=0\)) therefore returns \(+\infty\), not \(\mu\). A nonzero \(\gamma_i\) needs \(E[Y]<\infty\).

Return type:

Array

cov()[source]#

\(\mathrm{Cov}[X] = E[Y]\,\Sigma + \mathrm{Var}[Y]\,\gamma\gamma^\top\).

When \(\gamma=0\) only \(E[Y]\) is required (NInvG: \(\alpha>1\)). A zero entry of \(\gamma\gamma^\top\) contributes nothing even if \(\mathrm{Var}[Y]\) diverges. Non-existence is \(+\infty\).

Return type:

Array

skewness()[source]#

Component-wise skewness \(\gamma_1 = \mu_3 / \sigma^3\) of \(X\).

Closed form from the normal variance-mean mixture representation and subordinator raw moments; see The Generalized Hyperbolic Distribution. Shape (d,). Requires a finite third moment of the subordinator when \(\gamma \ne 0\). Symmetric coordinates return 0.

Return type:

Array

kurtosis()[source]#

Component-wise excess kurtosis \(\gamma_2 = \mu_4/\sigma^4 - 3\) of \(X\).

Closed form from the normal variance-mean mixture representation and subordinator raw moments; see The Generalized Hyperbolic Distribution. Shape (d,). Symmetric coordinates need \(E[Y^2]<\infty\) (NInvG: \(\alpha>2\)); a nonzero \(\gamma_i\) needs \(E[Y^4]\). Otherwise \(+\infty\) (e.g. \(t_3\)).

Return type:

Array

joint_entropy()[source]#

Entropy of the joint \((X, Y)\), \(H(X,Y) = C_\Sigma + H(Y) + \tfrac{d}{2}E[\log Y] + \tfrac{d}{2}\).

Return type:

Array

joint_varentropy()[source]#

Varentropy of the joint \((X, Y)\).

\[V_H(X, Y) = \tfrac{d}{2} + \mathrm{Var}\!\big[-\log g(Y) + \tfrac{d}{2}\log Y\big].\]

The marginal varentropy \(\mathrm{Var}[-\log f(X)]\) has no closed form; this tractable joint quantity is the natural surrogate.

Return type:

Array

joint_renyi(alpha)[source]#

Rényi entropy of order \(\alpha\) of the joint \((X, Y)\).

Parameters:

alpha (Array)

Return type:

Array

rvs(n, seed=42)[source]#

Sample X from the marginal distribution.

Parameters:
Return type:

Array

squared_hellinger(other)[source]#

Squared Hellinger distance via joint distributions (upper bound on marginal).

Parameters:

other (NormalMixture)

Return type:

Array

kl_divergence(other)[source]#

KL divergence via joint distributions.

Parameters:

other (NormalMixture)

Return type:

Array

compute_eta_from_model()[source]#

Reconstruct \(\eta\) from the model’s own parameters.

Uses the marginal expectations of the joint sufficient statistics:

\[\eta_4 = \mu + \gamma\,E[Y], \quad \eta_5 = \mu\,E[1/Y] + \gamma, \quad \eta_6 = \Sigma + \mu\mu^\top E[1/Y] + \gamma\gamma^\top E[Y] + \mu\gamma^\top + \gamma\mu^\top\]
Return type:

NormalMixtureEta

e_step(X, backend='jax')[source]#

Full E-step: subordinator conditionals + batch aggregation.

Returns EStepResult with aggregated NormalMixtureEta and the mean log-likelihood \(\ell_n(\theta)\) at the current parameters (nats per observation), via conjugacy \(\log f(x) = \log h(x) + \psi_{\mathrm{post}} - \psi\).

Parameters:
  • X ((n, d) data array)

  • backend (str) – 'jax' (default): jax.vmap over conditional_expectations. 'cpu': quad forms in JAX + GIG Bessel on CPU.

Return type:

EStepResult

classmethod from_expectation(eta, **kwargs)[source]#

Construct from expectation parameters \(\eta\).

Wraps JointNormalMixture.from_expectation(), which performs the exact closed-form M-step on the normal block and the subordinator’s from_expectation on the subordinator block.

This is the canonical η→model map: any prior or shrinkage target \(\eta_0\) can be inspected as a concrete model via cls.from_expectation(eta_0).sigma() etc.

Parameters:
  • eta (NormalMixtureEta) – Aggregated expectation parameters.

  • **kwargs – Forwarded to JointNormalMixture.from_expectation() (e.g. backend, method, maxiter, theta0).

Return type:

NormalMixture

m_step(eta, **kwargs)[source]#

Full M-step: update normal params + subordinator from \(\eta\).

Equivalent to type(self).from_expectation(eta, **kwargs); self is only used to dispatch on the subclass. Subclasses with iterative subordinator solvers (e.g. GeneralizedHyperbolic) may override to inject a warm-start \(\theta_0\).

Parameters:

eta (NormalMixtureEta)

Return type:

NormalMixture

m_step_normal(eta)[source]#

M-step for normal parameters only (MCECM Cycle 1).

Updates \(\mu, \gamma, \Sigma\); subordinator unchanged.

Parameters:

eta (NormalMixtureEta)

Return type:

NormalMixture

m_step_subordinator(eta, **kwargs)[source]#

M-step for the subordinator only (MCECM Cycle 2).

Reads the subordinator-relevant fields of eta; normal parameters are read from self._joint and copied unchanged. Subclasses with iterative solvers may override to add warm-start or sanity-check fallbacks.

Parameters:

eta (NormalMixtureEta)

Return type:

NormalMixture

replace(**updates)[source]#

Return a new model with selected parameters replaced.

Accepts any subset of:

  • normal parameters: mu, gamma, L_Sigma;

  • dispersion alias sigma — converted to L_Sigma via Cholesky (mutually exclusive with L_Sigma);

  • subordinator parameters declared by _subordinator_keys() (e.g. alpha, beta for VG / NInvG, mu_ig, lam for NIG, p, a, b for GH).

The actual storage lives in joint; this method does an immutable update via equinox.tree_at().

Examples

>>> vg2 = vg.replace(mu=new_mu)                    # change μ
>>> vg3 = vg.replace(alpha=2.5, beta=0.5)          # change subordinator
>>> vg4 = vg.replace(sigma=sigma2 * jnp.eye(d))    # set Σ via covariance
Return type:

NormalMixture

regularize_det_sigma(target_log_det=0.0)[source]#

Rescale to enforce \(\log|\Sigma| = \mathrm{target\_log\_det}\).

Picks \(s = \exp((\log|\Sigma| - \tau)/d)\) and applies _rescale(). The default target_log_det = 0 recovers the \(|\Sigma| = 1\) convention; passing the log-determinant of an initial reference Σ implements the det_sigma_x family.

Parameters:

target_log_det (float)

Return type:

NormalMixture

regularize_det_sigma_one()[source]#

Enforce \(|\Sigma| = 1\). Alias for regularize_det_sigma() with target_log_det = 0.

Return type:

NormalMixture

regularize_a_eq_b()[source]#

Rescale subordinator so that \(a = b = \sqrt{ab}\) for GIG-parameterised families.

Default implementation is a no-op; override in subclasses with both \(a, b > 0\) (currently GH and NIG; VG and NInvG have a degenerate a=0 or b=0 and the default no-op is the right behaviour).

Return type:

NormalMixture

em_convergence_params()[source]#

Diagnostic pytree (mu, gamma, L_Sigma) for hybrid-RMS param_changes.

Stopping uses the Aitken remaining gap of \(\ell_n\), not this pytree. Subordinator parameters are excluded here.

classmethod default_init(X)[source]#

Moment-based initialisation from data.

Returns a model with:

mu = sample mean gamma = zeros Sigma = empirical covariance (regularized) subordinator = distribution-specific defaults

Useful as a starting point for model.fit(X).

Parameters:

X (Array)

Return type:

NormalMixture

FactorNormalMixture#

Factor-analysis mixture base class: \(\Sigma = F F^\top + \mathrm{diag}(D)\) storage.

A factor-analysis Generalized-Hyperbolic random vector

\[X \stackrel{d}{=} \mu + \gamma Y + \sqrt{Y}\,(F Z + \varepsilon), \qquad Z \sim N(0, I_r),\; \varepsilon \sim N(0, \mathrm{diag}(D)),\]

where \(Y\) is a GH subordinator (Gamma, InverseGamma, InverseGaussian or GIG). The marginal \(X\) is a normal variance-mean mixture with dispersion \(\Sigma = F F^\top + \mathrm{diag}(D)\). Storing \((F, D)\) instead of a full Cholesky of \(\Sigma\) keeps quadratic forms and log-determinants in \(\mathcal{O}(d r^2 + r^3)\) via Woodbury and makes the rotation gauge of \(F\) irrelevant for the hybrid-RMS diagnostic (we measure on \(\Sigma\)).

This module hosts only the abstract FactorNormalMixture base. The four concrete subordinator families (FactorVarianceGamma, FactorNormalInverseGamma, FactorNormalInverseGaussian, FactorGeneralizedHyperbolic) live next to their full-\(\Sigma\) siblings in normix/distributions/.

The FA complete-data structure is over \((X, Y, Z)\) with ten sufficient statistics (see docs/theory/factor_analysis.md), so this family does not share the joint’s six-statistic exponential-family signature — that is why FactorNormalMixture is a sibling of NormalMixture, not a subclass.

class normix.mixtures.factor.FactorNormalMixture(mu, gamma, F, D, _subordinator)[source]#

Bases: MarginalMixture

Abstract marginal \(f(x)\) for a factor-analysis normal variance-mean mixture with dispersion \(\Sigma = F F^\top + \mathrm{diag}(D)\).

Stored fields#

mu(d,)

Location \(\mu\).

gamma(d,)

Skewness \(\gamma\).

F(d, r)

Factor loadings; r is the latent factor dimension.

D(d,)

Diagonal entries of the residual covariance (positive).

_subordinatorExponentialFamily

Fitted subordinator (Gamma / InverseGamma / InverseGaussian / GIG). Access via subordinator().

Subclasses define the subordinator family (via the instance stored in _subordinator, which supplies the shared posterior map through to_gig()) and implement log_prob(), _subordinator_from_eta(), and a few forwarders / initialisers. The prior-to-posterior GIG conjugacy (_posterior_gig_params()) is uniform across families and lives on the base. The linear algebra for \(\Sigma^{-1}\) and \(\log|\Sigma|\) is shared via Woodbury helpers (_M(), _solve(), _quad_form(), _log_det_sigma(), _beta()).

Notes

F is identifiable only up to a right \(r \times r\) orthogonal rotation, so the hybrid-RMS diagnostic is on \(\Sigma = F F^\top + \mathrm{diag}(D)\) (em_convergence_params()) rather than on F directly. Stopping uses the Aitken remaining gap of \(\ell_n\).

mu: Array#
gamma: Array#
F: Array#
D: Array#
property d: int#
property r: int#
subordinator()[source]#

Return the fitted subordinator distribution \(Y\).

Return type:

ExponentialFamily

sigma()[source]#

\(\Sigma = F F^\top + \mathrm{diag}(D)\) (dense).

Return type:

Array

log_det_sigma()[source]#

\(\log|\Sigma|\), computed via Woodbury.

Return type:

Array

mean()[source]#

\(E[X] = \mu + \gamma\,E[Y]\) (see NormalMixture.mean()).

Return type:

Array

cov()[source]#

\(\mathrm{Cov}[X] = E[Y]\,\Sigma + \mathrm{Var}[Y]\,\gamma\gamma^\top\) (see NormalMixture.cov()).

Return type:

Array

skewness()[source]#

Component-wise skewness of \(X\) (see NormalMixture.skewness()).

Return type:

Array

kurtosis()[source]#

Component-wise excess kurtosis of \(X\) (see NormalMixture.kurtosis()).

Return type:

Array

rvs(n, seed=42)[source]#

Sample n observations from the marginal \(f(x)\).

Uses \(X = \mu + \gamma Y + \sqrt{Y}(F Z + \varepsilon)\).

Parameters:
Return type:

Array

e_step(X, *, backend='jax')[source]#

Full E-step: posterior \(Y\) expectations + deterministic \(Z\) reductions.

Returns EStepResult whose eta is a FactorMixtureStats (first six fields are batch averages of the standard mixture statistics; the four \(Z\)-fields follow docs/theory/factor_analysis.md §E-Step) and whose log_lik is the conjugacy mean log-likelihood at the current parameters.

Parameters:
Return type:

EStepResult

m_step(eta, **kwargs)[source]#

Full M-step: update \((\mu, \gamma, F, D)\) and the subordinator from FactorMixtureStats.

Equivalent to self.m_step_normal(eta).m_step_subordinator(eta, **kwargs). Subclasses with iterative subordinator solvers (e.g. FactorGeneralizedHyperbolic) may override to inject a warm-start.

Parameters:

eta (FactorMixtureStats)

Return type:

FactorNormalMixture

m_step_normal(eta)[source]#

M-step for \((\mu, \gamma, F, D)\) only (MCECM cycle 1).

Subordinator unchanged.

Parameters:

eta (FactorMixtureStats)

Return type:

FactorNormalMixture

m_step_subordinator(eta, **kwargs)[source]#

M-step for the subordinator only (MCECM cycle 2).

Reads the subordinator-relevant fields of eta; \((\mu, \gamma, F, D)\) are read from self and copied unchanged. Subclasses with iterative solvers may override to add a warm-start or sanity-check fallback.

Parameters:

eta (FactorMixtureStats)

Return type:

FactorNormalMixture

compute_eta_from_model()[source]#

Reconstruct FactorMixtureStats from the model’s own parameters.

The first six fields use the marginal expectations of the joint sufficient statistics under the current \(\Sigma = F F^\top + \mathrm{diag}(D)\); the four \(Z\)-fields use the same deterministic relations as e_step() (§E-Step of the theory doc).

Return type:

FactorMixtureStats

em_convergence_params()[source]#

Diagnostic pytree \((\mu, \gamma, \Sigma)\) for hybrid-RMS param_changes.

\(\Sigma = F F^\top + \mathrm{diag}(D)\) rather than F directly, because F is identifiable only up to an \(r \times r\) orthogonal rotation. Stopping uses the Aitken remaining gap of \(\ell_n\), not this pytree.

regularize_det_sigma(target_log_det=0.0)[source]#

Rescale to enforce \(\log|\Sigma| = \mathrm{target\_log\_det}\).

Same family as the full-Σ implementation but the log-determinant is computed via the Woodbury identity \(\log|\Sigma| = \log|D| + \log|I_r + F^\top D^{-1} F|\).

Parameters:

target_log_det (float)

Return type:

FactorNormalMixture

regularize_det_sigma_one()[source]#

Enforce \(|\Sigma| = 1\). Alias for regularize_det_sigma() with target_log_det = 0.

Return type:

FactorNormalMixture

regularize_a_eq_b()[source]#

Rescale subordinator so that \(a = b = \sqrt{ab}\).

Default no-op; overridden in FactorGeneralizedHyperbolic and FactorNormalInverseGaussian.

Return type:

FactorNormalMixture

classmethod default_init(X, *, r=1)[source]#

Moment-based initialisation from data with r factors.

Sets \(\mu = \bar X\), \(\gamma = 0\), and splits the empirical covariance into a rank-r factor block (top-r eigenvectors scaled by \(\sqrt{\lambda - \bar\lambda}\)) plus a positive diagonal residual.

Useful as a starting point for model.fit(X).

Parameters:
Return type:

FactorNormalMixture

replace(**updates)[source]#

Return a new model with selected top-level fields replaced.

Accepts any subset of _NORMAL_KEYS. The public key subordinator updates the private _subordinator field. Subclass-specific subordinator parameter shortcuts (e.g. alpha=...) are not supported on the FA family — use replace(subordinator=Gamma(alpha=..., beta=...)) instead. This keeps the contract narrow (one storage form per family).

Return type:

FactorNormalMixture

Parameters: