Vectorized Gaussian Mixture Model for adaptive importance sampling.

Stores all mixture components as numpy arrays and provides vectorized sampling, PDF evaluation, and expectation-maximization (EM) updates.

class cosmic.sample.stroopwafel.mixture_model.GaussianMixture(means, covariances, alphas, rejection_rate=0.0, min_active_fraction=0.01, min_entropy_change=0.01)

Bases: object

A mixture of K multivariate Gaussians.

Parameters:

meansnumpy.ndarray

(K, D) array of component means.

covariancesnumpy.ndarray

(K, D, D) array of covariance matrices.

alphasnumpy.ndarray

(K,) array of mixture weights (must sum to 1).

rejection_ratefloat, optional

Fraction of samples that fall outside bounds, by default 0.0

component_pdfs(samples)

Evaluate each component’s PDF at the given samples.

Parameters:

samplesnumpy.ndarray

(N, D) array in sampling space.

Returns:

numpy.ndarray

(K, N) array where element [k, n] is the PDF of component k evaluated at sample n.

compute_rejection_rate(param_space, reject_mask_fn, n_per_component=10000, rng=None)

Estimate the rejection rate of the mixture.

Samples from each component, transforms to physical space, applies rejection criteria, and computes the weighted rejection rate.

Parameters:

param_spaceParameterSpace

Parameter space for bounds checking and coordinate transforms.

reject_mask_fncallable

Function (samples_physical) -> bool_mask returning True for physically rejected systems.

n_per_componentint, optional

Number of samples per component for the estimate, by default 10000

rngnumpy.random.Generator, optional

Random number generator, by default None

Returns:

float

Estimated rejection rate (also stored as self.rejection_rate).

classmethod from_hits(hit_samples, param_space, average_density_one_dim, kappa=1.0, min_active_fraction=0.01, min_entropy_change=0.01)

Create a Gaussian mixture by placing one component at each hit.

Parameters:

hit_samplesnumpy.ndarray

(K, D) array of hit locations in sampling space.

param_spaceParameterSpace

Parameter space instance providing bounds and sigma computation.

average_density_one_dimfloat

Characteristic inter-sample spacing, 1 / num_explored ** (1 / D).

kappafloat, optional

Width scaling factor for the Gaussian covariances, by default 1.0

min_active_fractionfloat, optional

Minimum value of (1 - rejection_rate) used in oversampling and normalisation calculations, by default 0.01

min_entropy_changefloat, optional

Minimum change in normalised effective sample size (entropy) to avoid reverting to the previous mixture state, by default 0.01

Returns:

GaussianMixture

A new mixture with one component centred on each hit.

property n_components

property ndim

pdf(samples)

Evaluate the mixture PDF at the given samples.

Parameters:

samplesnumpy.ndarray

(N, D) array in sampling space.

Returns:

numpy.ndarray

(N,) array of PDF values (weighted sum of components).

sample(n_total, param_space, consider_rejection=False, rng=None)

Sample from the mixture distribution.

Parameters:

n_totalint

Total number of samples desired.

param_spaceParameterSpace

Parameter space used for bounds checking.

consider_rejectionbool, optional

If True, oversample to account for the current rejection rate, by default False

rngnumpy.random.Generator, optional

Random number generator, by default None

Returns:

samplesnumpy.ndarray

(M, D) array of samples in sampling space.

masknumpy.ndarray

(M,) boolean array indicating which samples are in bounds.

gaussian_indicesnumpy.ndarray

(M,) integer array indicating which component generated each sample.

update_em(samples, is_hit, prior_probs, prior_fraction_rejected, tolerance=1e-10, entropies=None)

Perform one (importance-weighted) expectation-maximization (EM) update.

EM is the standard algorithm for fitting a mixture model. It alternates an E-step – computing each component’s responsibility for every sample (the posterior probability that the sample was drawn from that component) – with an M-step that re-estimates every component’s weight, mean, and covariance as the responsibility-weighted moments of the samples. Here the samples additionally carry importance weights (prior_probs * is_hit / q), so productive components (those near many hits) gain weight and recentre on where the hits actually are; components whose weight falls below tolerance are dropped.

The update is only worth keeping if it improves the proposal. This is measured by the normalised effective sample size exp(H(w)) / N (with H the Shannon entropy of the normalised weights): if it fails to increase by at least min_entropy_change over the previous generation, the method signals a revert (see Returns) – STROOPWAFEL’s indication that the mixture has stopped improving.

Parameters:

samplesnumpy.ndarray

(N, D) array of samples in sampling space.

is_hitnumpy.ndarray

(N,) boolean or integer array indicating hits.

prior_probsnumpy.ndarray

(N,) array of prior probabilities for each sample.

prior_fraction_rejectedfloat

Fraction of prior samples that are physically rejected.

tolerancefloat, optional

Minimum mixture weight to keep a component, by default 1e-10

entropieslist, optional

List of previous entropy values (mutated in place for convergence tracking), by default None

Returns:

bool

True if the entropy check triggers reversion to the previous mixture state.