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)\):
Sufficient statistics:
Natural parameters:
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:
Expectation parameters (EM E-step quantities):
EM M-step closed-form (let \(D = 1 - E[1/Y] \cdot E[Y]\)):
- class normix.mixtures.joint.JointNormalMixture(mu, gamma, L_Sigma)[source]#
Bases:
ExponentialFamilyAbstract 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.- 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 = 0stores the Gamma/InvGamma boundary exactly (required for divergence gauges).- Parameters:
boundary_eps (float)
- log_det_sigma()[source]#
\(\log|\Sigma| = 2\sum_i \log L_{ii}\), via Cholesky diagonal.
- Return type:
- 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.
- 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 keysE_log_Y,E_inv_Y,E_Y, pluspsi_post(posterior GIG log-partition) andzw(\(\gamma^\top\Sigma^{-1}(x-\mu)\)) for the conjugacy identity \(\log f(x) = \log h(x) + \psi_{\mathrm{post}} - \psi\).
- 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]\).
- 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 fromExponentialFamily.
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,theta0for warm-starting GIG). For the flat-array path: forwarded to the parent solver.
- Return type:
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-densitye_step(X)—jax.vmap()over conditional expectationsm_step(X, expectations)— returns newNormalMixturefit(X, ...)— convenience EM fitting with multi-start
- class normix.mixtures.marginal.MarginalMixture[source]#
Bases:
ModuleAbstract 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 inFactorNormalMixture) and the type of the EM expectation pytree (NormalMixtureEtavs.FactorMixtureStats).The fitter depends only on this contract; it does not know which storage form a model uses.
- abstractmethod e_step(X, *, backend='jax')[source]#
E-step:
EStepResult(eta, log_lik).
- abstractmethod m_step(eta, **kwargs)[source]#
Full M-step: updates all parameters; returns a new model.
- Parameters:
eta (Any)
- Return type:
- abstractmethod m_step_normal(eta)[source]#
M-step for normal parameters only (MCECM cycle 1).
- Parameters:
eta (Any)
- Return type:
- abstractmethod m_step_subordinator(eta, **kwargs)[source]#
M-step for subordinator parameters only (MCECM cycle 2).
- Parameters:
eta (Any)
- Return type:
- abstractmethod compute_eta_from_model()[source]#
Reconstruct the expectation pytree from the model’s own parameters.
- Return type:
- 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_changesonly. 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:
- fit(X, *, alpha_min=None, **fitter_kwargs)[source]#
Fit using
selfas initialisation. ReturnsEMResult.Keyword arguments other than
alpha_minare forwarded toBatchEMFitter(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 intom_step_kwargsand wins over anyalpha_minalready present there.X (jax.Array)
- Return type:
- class normix.mixtures.marginal.NormalMixture(joint)[source]#
Bases:
MarginalMixtureMarginal \(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; usereplace()to obtain a new model with updated parameters (modules are immutable).- property joint: JointNormalMixture#
- 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:
- 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:
- 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:
- 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:
- 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:
- 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:
- 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:
- squared_hellinger(other)[source]#
Squared Hellinger distance via joint distributions (upper bound on marginal).
- Parameters:
other (NormalMixture)
- Return type:
- kl_divergence(other)[source]#
KL divergence via joint distributions.
- Parameters:
other (NormalMixture)
- Return type:
- 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:
- e_step(X, backend='jax')[source]#
Full E-step: subordinator conditionals + batch aggregation.
Returns
EStepResultwith aggregatedNormalMixtureEtaand 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:
- Return type:
- 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’sfrom_expectationon 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:
- m_step(eta, **kwargs)[source]#
Full M-step: update normal params + subordinator from \(\eta\).
Equivalent to
type(self).from_expectation(eta, **kwargs);selfis 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:
- 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:
- 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 fromself._jointand copied unchanged. Subclasses with iterative solvers may override to add warm-start or sanity-check fallbacks.- Parameters:
eta (NormalMixtureEta)
- Return type:
- 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 toL_Sigmavia Cholesky (mutually exclusive withL_Sigma);subordinator parameters declared by
_subordinator_keys()(e.g.alpha, betafor VG / NInvG,mu_ig, lamfor NIG,p, a, bfor GH).
The actual storage lives in
joint; this method does an immutable update viaequinox.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:
- 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 defaulttarget_log_det = 0recovers the \(|\Sigma| = 1\) convention; passing the log-determinant of an initial reference Σ implements thedet_sigma_xfamily.- Parameters:
target_log_det (float)
- Return type:
- regularize_det_sigma_one()[source]#
Enforce \(|\Sigma| = 1\). Alias for
regularize_det_sigma()withtarget_log_det = 0.- Return type:
- 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=0orb=0and the default no-op is the right behaviour).- Return type:
- em_convergence_params()[source]#
Diagnostic pytree
(mu, gamma, L_Sigma)for hybrid-RMSparam_changes.Stopping uses the Aitken remaining gap of \(\ell_n\), not this pytree. Subordinator parameters are excluded here.
FactorNormalMixture#
Factor-analysis mixture base class: \(\Sigma = F F^\top + \mathrm{diag}(D)\) storage.
A factor-analysis Generalized-Hyperbolic random vector
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:
MarginalMixtureAbstract 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;
ris 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 throughto_gig()) and implementlog_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
Fis 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 onFdirectly. Stopping uses the Aitken remaining gap of \(\ell_n\).- cov()[source]#
\(\mathrm{Cov}[X] = E[Y]\,\Sigma + \mathrm{Var}[Y]\,\gamma\gamma^\top\) (see
NormalMixture.cov()).- Return type:
- kurtosis()[source]#
Component-wise excess kurtosis of \(X\) (see
NormalMixture.kurtosis()).- Return type:
- rvs(n, seed=42)[source]#
Sample
nobservations from the marginal \(f(x)\).Uses \(X = \mu + \gamma Y + \sqrt{Y}(F Z + \varepsilon)\).
- e_step(X, *, backend='jax')[source]#
Full E-step: posterior \(Y\) expectations + deterministic \(Z\) reductions.
Returns
EStepResultwhoseetais aFactorMixtureStats(first six fields are batch averages of the standard mixture statistics; the four \(Z\)-fields followdocs/theory/factor_analysis.md§E-Step) and whoselog_likis the conjugacy mean log-likelihood at the current parameters.- Parameters:
- Return type:
- 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:
- m_step_normal(eta)[source]#
M-step for \((\mu, \gamma, F, D)\) only (MCECM cycle 1).
Subordinator unchanged.
- Parameters:
eta (FactorMixtureStats)
- Return type:
- 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 fromselfand copied unchanged. Subclasses with iterative solvers may override to add a warm-start or sanity-check fallback.- Parameters:
eta (FactorMixtureStats)
- Return type:
- compute_eta_from_model()[source]#
Reconstruct
FactorMixtureStatsfrom 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:
- em_convergence_params()[source]#
Diagnostic pytree \((\mu, \gamma, \Sigma)\) for hybrid-RMS
param_changes.\(\Sigma = F F^\top + \mathrm{diag}(D)\) rather than
Fdirectly, becauseFis 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:
- regularize_det_sigma_one()[source]#
Enforce \(|\Sigma| = 1\). Alias for
regularize_det_sigma()withtarget_log_det = 0.- Return type:
- regularize_a_eq_b()[source]#
Rescale subordinator so that \(a = b = \sqrt{ab}\).
Default no-op; overridden in
FactorGeneralizedHyperbolicandFactorNormalInverseGaussian.- Return type:
- classmethod default_init(X, *, r=1)[source]#
Moment-based initialisation from data with
rfactors.Sets \(\mu = \bar X\), \(\gamma = 0\), and splits the empirical covariance into a rank-
rfactor block (top-reigenvectors scaled by \(\sqrt{\lambda - \bar\lambda}\)) plus a positive diagonal residual.Useful as a starting point for
model.fit(X).- Parameters:
- Return type:
- replace(**updates)[source]#
Return a new model with selected top-level fields replaced.
Accepts any subset of
_NORMAL_KEYS. The public keysubordinatorupdates the private_subordinatorfield. Subclass-specific subordinator parameter shortcuts (e.g.alpha=...) are not supported on the FA family — usereplace(subordinator=Gamma(alpha=..., beta=...))instead. This keeps the contract narrow (one storage form per family).- Return type:
- Parameters:
mu (Array)
gamma (Array)
F (Array)
D (Array)
_subordinator (ExponentialFamily)