mixle.models.mixture_density module

NeuralConditionalDensity – the adapter that turns ANY torch conditional density into a mixle leaf.

This is the conditional sibling of NeuralDensity. Where that one wraps a module exposing log_density(x) -> (n,) (an unconditional p(x)), this wraps a module exposing log_density(x, y) -> (n,) (and sample_given(x) -> (n, d)) and gives you a full five-piece mixle Distribution over the pair (x, y) – so a flexible conditional density drops into a mixture of experts, a composite field, or an HMM emission and is fit jointly with classical families by the same responsibility-weighted-NLL EM M-step (warm-started across iterations, i.e. generalized EM).

Why it matters: NeuralGaussian fixes the conditional law to a single Gaussian, p(y | x) = N(y; f(x), sigma^2 I) – one mean per x, unimodal and homoscedastic. Many real conditionals are neither: an inverse problem has several valid y for one x; measurement noise grows with x. build_mdn() is the ready instance – a mixture density network, p(y | x) = sum_k pi_k(x) N(y; mu_k(x), sigma_k(x)^2) – whose entire mixture (weights, means, variances) is a function of x, so it is multimodal and heteroscedastic. Any other conditional density (a conditional flow, an autoregressive head) plugs in the same way: give it log_density(x, y) and sample_given(x).

class NeuralConditionalDensity(module, *, m_steps=60, lr=5e-3, device='cpu', name=None)

Bases: SequenceEncodableProbabilityDistribution

Wrap a torch conditional-density module (module.log_density(x, y) -> (n,)) as a mixle leaf.

Observations are pairs (x, y). The module must also expose sample_given(x) -> (n, d) to draw y.

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• device (str)

• name (str | None)

log_density(xy)

Return the log-density or log-mass at a single observation.

Parameters:

xy (Any)

Return type:

float

seq_log_density(enc)

Return vectorized log-density values for sequence-encoded observations.

Parameters:

enc (Any)

Return type:

ndarray

sampler(seed=None)

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

NeuralConditionalDensitySampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

NeuralConditionalDensityEstimator

dist_to_encoder()

Return the data encoder used by this distribution for vectorized methods.

Return type:

NeuralConditionalDensityEncoder

to_dict()

Return a safe JSON-compatible representation of this distribution.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Reconstruct a distribution from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

NeuralConditionalDensity

class NeuralConditionalDensitySampler(dist, seed=None)

Bases: DistributionSampler

Parameters:

• dist (NeuralConditionalDensity)

• seed (int | None)

sample(size=None, *, batched=True)

Draw observations.

Combinator samplers (mixture/sequence/…) accept batched. With batched=True (the default) each child stream is drawn in one vectorized call instead of a per-draw Python loop – far faster. Because every child sampler owns an independent RandomState, batching consumes each stream in the same order as the loop, so the draws are identical to the legacy path. batched=False forces that legacy per-draw loop as a guaranteed- stable reference. Leaf samplers are already vectorized and ignore the flag.

Parameters:

• size (int | None)

• batched (bool)

Return type:

Any

sample_given(x)

Parameters:

x (Any)

Return type:

ndarray

class NeuralConditionalDensityEncoder

Bases: DataSequenceEncoder

seq_encode(data)

Encode the iid observation sequence x for vectorized evaluation.

Parameters:

data (list)

Return type:

tuple[ndarray, ndarray]

class NeuralConditionalDensityAccumulator

Bases: SequenceEncodableStatisticAccumulator

Buffers responsibility-weighted (x, y) pairs for the M-step (the weights are the E-step soft counts).

update(xy, weight, estimate)

Parameters:

• xy (Any)

• weight (float)

• estimate (Any)

Return type:

None

seq_update(enc, weights, estimate)

Parameters:

• enc (Any)

• weights (ndarray)

• estimate (Any)

Return type:

None

initialize(xy, weight, rng)

Parameters:

• xy (Any)

• weight (float)

• rng (Any)

Return type:

None

seq_initialize(enc, weights, rng)

Parameters:

• enc (Any)

• weights (ndarray)

• rng (Any)

Return type:

None

combine(other)

Parameters:

other (Any)

Return type:

NeuralConditionalDensityAccumulator

value()

Return type:

tuple

from_value(value)

Parameters:

value (tuple)

Return type:

NeuralConditionalDensityAccumulator

acc_to_encoder()

Return type:

NeuralConditionalDensityEncoder

class NeuralConditionalDensityAccumulatorFactory

Bases: StatisticAccumulatorFactory

make()

Return type:

NeuralConditionalDensityAccumulator

class NeuralConditionalDensityEstimator(module, *, m_steps=60, lr=5e-3, device='cpu', name=None)

Bases: ParameterEstimator

M-step: responsibility-weighted MLE max sum_i w_i log p(y_i | x_i) by gradient ascent (warm-started).

Parameters:

• module (Any)

• m_steps (int)

• lr (float)

• device (str)

• name (str | None)

accumulator_factory()

Return type:

NeuralConditionalDensityAccumulatorFactory

estimate(nobs, suff_stat)

Parameters:

• nobs (float | None)

• suff_stat (tuple)

Return type:

NeuralConditionalDensity

build_mdn(x_dim, y_dim, *, k=5, hidden=32, layers=2)

A mixture density network: p(y | x) = sum_k pi_k(x) N(y; mu_k(x), diag sigma_k(x)^2) – ready to wrap.

A shared MLP body maps x to three heads – mixing logits, component means, and (log) component scales – so the entire conditional law is a function of x: multimodal (several mu_k) and heteroscedastic (input-dependent sigma_k). Exposes log_density(x, y) (a log-sum-exp over components) and sample_given(x) (pick a component by pi, then a Gaussian), the contract a NeuralConditionalDensity adapts.

Parameters:

• x_dim (int)

• y_dim (int)

• k (int)

• hidden (int)

• layers (int)

Return type:

Any

build_conditional_flow(x_dim, y_dim, *, hidden=32, layers=4)

A conditional coupling flow: an exact p(y | x) whose transform of y is conditioned on x.

The exact-density counterpart to build_mdn(). Each affine-coupling layer’s shift/scale networks take both the passed-through y coordinates and x, so the whole invertible y-transform bends with the input – capturing within-``y`` dependence (e.g. y2 a nonlinear function of y1) that a single-Gaussian NeuralGaussian (isotropic mean-only) cannot, while keeping an exact log-density (so it composes honestly, unlike a bound). Needs y_dim >= 2 for the coupling to be non-trivial. Exposes log_density(x, y) and sample_given(x) – the contract a NeuralConditionalDensity adapts.

Parameters:

• x_dim (int)

• y_dim (int)

• hidden (int)

• layers (int)

Return type:

Any

build_conditional_autoregressive_categorical(x_dim, y_dim, n_categories, *, hidden=64)

An autoregressive categorical conditioned on x: exact p(y | x) over discrete y in {0..C-1}^y_dim.

The conditional sibling of build_autoregressive_categorical() and the discrete counterpart to build_conditional_flow(). It factorizes p(y | x) = prod_i p(y_i | y_{<i}, x) with a MADE-masked net over y into which x is injected unmasked (degree 0, so every coordinate may depend on x). Each per-coordinate softmax is exactly a conditional, so the density is exactly normalized and composes honestly. Exposes log_density(x, y) and sample_given(x) – the contract a NeuralConditionalDensity adapts.

Parameters:

• x_dim (int)

• y_dim (int)

• n_categories (int)

• hidden (int)

Return type:

Any