def test_mixture(
distr1: Distribution, distr2: Distribution, p: Tensor
) -> None:
# sample from component distributions, and select samples
samples1 = distr1.sample(num_samples=NUM_SAMPLES_LARGE)
samples2 = distr2.sample(num_samples=NUM_SAMPLES_LARGE)
rand = mx.nd.random.uniform(shape=(NUM_SAMPLES_LARGE, *p.shape))
choice = (rand < p.expand_dims(axis=0)).broadcast_like(samples1)
samples_ref = mx.nd.where(choice, samples1, samples2)
# construct mixture distribution and sample from it
mixture_probs = mx.nd.stack(p, 1.0 - p, axis=-1)
mixture = MixtureDistribution(
mixture_probs=mixture_probs, components=[distr1, distr2]
)
samples_mix = mixture.sample(num_samples=NUM_SAMPLES_LARGE)
# check that shapes are right
assert (
samples1.shape
== samples2.shape
== samples_mix.shape
== samples_ref.shape
)
# check mean and stddev
calc_mean = mixture.mean.asnumpy()
sample_mean = samples_mix.asnumpy().mean(axis=0)
assert np.allclose(calc_mean, sample_mean, atol=1e-1)
# check that histograms are close
assert (
diff(
histogram(samples_mix.asnumpy()), histogram(samples_ref.asnumpy())
)
< 0.05
)
# can only calculated cdf for gaussians currently
if isinstance(distr1, Gaussian) and isinstance(distr2, Gaussian):
emp_cdf, edges = empirical_cdf(samples_mix.asnumpy())
calc_cdf = mixture.cdf(mx.nd.array(edges)).asnumpy()
assert np.allclose(calc_cdf[1:, :], emp_cdf, atol=1e-2)
def test_mixture(
distr1: Distribution, distr2: Distribution, p: Tensor
) -> None:
# sample from component distributions, and select samples
samples1 = distr1.sample(num_samples=NUM_SAMPLES)
samples2 = distr2.sample(num_samples=NUM_SAMPLES)
rand = mx.nd.random.uniform(shape=(NUM_SAMPLES, *p.shape))
choice = (rand < p.expand_dims(axis=0)).broadcast_like(samples1)
samples_ref = mx.nd.where(choice, samples1, samples2)
# construct mixture distribution and sample from it
mixture_probs = mx.nd.stack(p, 1.0 - p, axis=-1)
mixture = MixtureDistribution(
mixture_probs=mixture_probs, components=[distr1, distr2]
)
samples_mix = mixture.sample(num_samples=NUM_SAMPLES)
# check that shapes are right
assert (
samples1.shape
== samples2.shape
== samples_mix.shape
== samples_ref.shape
)
# check that histograms are close
assert (
diff(
histogram(samples_mix.asnumpy()), histogram(samples_ref.asnumpy())
)
< 0.05
)