def test_draw_samples_with_broadcast(self, dtype, mean, mean_isSamples,
var, var_isSamples, rv_shape,
num_samples):
mean_mx = mx.nd.array(mean, dtype=dtype)
if not mean_isSamples:
mean_mx = add_sample_dimension(mx.nd, mean_mx)
var_mx = mx.nd.array(var, dtype=dtype)
if not var_isSamples:
var_mx = add_sample_dimension(mx.nd, var_mx)
var = var_mx.asnumpy()
isSamples_any = any([mean_isSamples, var_isSamples])
rand = np.random.rand(num_samples, *rv_shape)
rand_gen = MockMXNetRandomGenerator(
mx.nd.array(rand.flatten(), dtype=dtype))
rv_samples_np = mean + np.matmul(np.linalg.cholesky(var),
np.expand_dims(rand, axis=-1)).sum(-1)
normal = MultivariateNormal.define_variable(shape=rv_shape,
dtype=dtype,
rand_gen=rand_gen).factor
variables = {normal.mean.uuid: mean_mx, normal.covariance.uuid: var_mx}
draw_samples_rt = normal.draw_samples(F=mx.nd, variables=variables)
assert np.issubdtype(draw_samples_rt.dtype, dtype)
assert is_sampled_array(mx.nd, draw_samples_rt) == isSamples_any
if isSamples_any:
assert get_num_samples(
mx.nd, draw_samples_rt) == num_samples, (get_num_samples(
mx.nd, draw_samples_rt), num_samples)
def test_log_pdf_no_broadcast(self, dtype, mean, mean_isSamples, var,
var_isSamples, rv, rv_isSamples,
num_samples):
mean_mx = mx.nd.array(mean, dtype=dtype)
if not mean_isSamples:
mean_mx = add_sample_dimension(mx.nd, mean_mx)
mean = mean_mx.asnumpy()
var_mx = mx.nd.array(var, dtype=dtype)
if not var_isSamples:
var_mx = add_sample_dimension(mx.nd, var_mx)
var = var_mx.asnumpy()
rv_mx = mx.nd.array(rv, dtype=dtype)
if not rv_isSamples:
rv_mx = add_sample_dimension(mx.nd, rv_mx)
rv = rv_mx.asnumpy()
from scipy.stats import multivariate_normal
isSamples_any = any([mean_isSamples, var_isSamples, rv_isSamples])
rv_shape = rv.shape[1:]
n_dim = 1 + len(
rv.shape) if isSamples_any and not rv_isSamples else len(rv.shape)
mean_np = numpy_array_reshape(mean, isSamples_any, n_dim)
var_np = numpy_array_reshape(var, isSamples_any, n_dim)
rv_np = numpy_array_reshape(rv, isSamples_any, n_dim)
rand = np.random.rand(num_samples, *rv_shape)
rand_gen = MockMXNetRandomGenerator(
mx.nd.array(rand.flatten(), dtype=dtype))
r = []
for s in range(len(rv_np)):
a = []
for i in range(len(rv_np[s])):
a.append(
multivariate_normal.logpdf(rv_np[s][i], mean_np[s][i],
var_np[s][i]))
r.append(a)
log_pdf_np = np.array(r)
normal = MultivariateNormal.define_variable(shape=rv_shape,
dtype=dtype,
rand_gen=rand_gen).factor
variables = {
normal.mean.uuid: mean_mx,
normal.covariance.uuid: var_mx,
normal.random_variable.uuid: rv_mx
}
log_pdf_rt = normal.log_pdf(F=mx.nd, variables=variables)
assert np.issubdtype(log_pdf_rt.dtype, dtype)
assert is_sampled_array(mx.nd, log_pdf_rt) == isSamples_any
if isSamples_any:
assert get_num_samples(
mx.nd, log_pdf_rt) == num_samples, (get_num_samples(
mx.nd, log_pdf_rt), num_samples)
assert np.allclose(log_pdf_np, log_pdf_rt.asnumpy())
def test_draw_samples_no_broadcast(self, dtype, mean, mean_isSamples, var,
var_isSamples, rv_shape, num_samples):
mean_mx = mx.nd.array(mean, dtype=dtype)
if not mean_isSamples:
mean_mx = add_sample_dimension(mx.nd, mean_mx)
var_mx = mx.nd.array(var, dtype=dtype)
if not var_isSamples:
var_mx = add_sample_dimension(mx.nd, var_mx)
var = var_mx.asnumpy()
# n_dim = 1 + len(rv.shape) if isSamples_any else len(rv.shape)
rand = np.random.rand(num_samples, *rv_shape)
rand_gen = MockMXNetRandomGenerator(
mx.nd.array(rand.flatten(), dtype=dtype))
rand_exp = np.expand_dims(rand, axis=-1)
lmat = np.linalg.cholesky(var)
temp1 = np.matmul(lmat, rand_exp).sum(-1)
rv_samples_np = mean + temp1
normal = MultivariateNormal.define_variable(shape=rv_shape,
dtype=dtype,
rand_gen=rand_gen).factor
variables = {normal.mean.uuid: mean_mx, normal.covariance.uuid: var_mx}
draw_samples_rt = normal.draw_samples(F=mx.nd,
variables=variables,
num_samples=num_samples)
assert np.issubdtype(draw_samples_rt.dtype, dtype)
assert array_has_samples(mx.nd, draw_samples_rt) == True
assert get_num_samples(
mx.nd, draw_samples_rt) == num_samples, (get_num_samples(
mx.nd, draw_samples_rt), num_samples)
def test_draw_samples_with_broadcast_no_numpy_verification(
self, dtype, mean, mean_isSamples, var, var_isSamples, rv_shape,
num_samples):
mean_mx = mx.nd.array(mean, dtype=dtype)
if not mean_isSamples:
mean_mx = add_sample_dimension(mx.nd, mean_mx)
var_mx = mx.nd.array(var, dtype=dtype)
if not var_isSamples:
var_mx = add_sample_dimension(mx.nd, var_mx)
var = var_mx.asnumpy()
rand = np.random.rand(num_samples, *rv_shape)
rand_gen = MockMXNetRandomGenerator(
mx.nd.array(rand.flatten(), dtype=dtype))
normal = MultivariateNormal.define_variable(shape=rv_shape,
dtype=dtype,
rand_gen=rand_gen).factor
variables = {normal.mean.uuid: mean_mx, normal.covariance.uuid: var_mx}
draw_samples_rt = normal.draw_samples(F=mx.nd,
variables=variables,
num_samples=num_samples)
assert np.issubdtype(draw_samples_rt.dtype, dtype)
assert is_sampled_array(mx.nd, draw_samples_rt) == True
def test_draw_samples_non_mock(self, plot=False):
# Also make sure the non-mock sampler works
dtype = np.float32
num_samples = 100000
mean = np.array([0.5, 0])
covariance = np.array([[2, 0.5], [0.5, 2]])
rv_shape = (2,)
mean_mx = add_sample_dimension(mx.nd, mx.nd.array(mean, dtype=dtype))
covariance_mx = add_sample_dimension(mx.nd, mx.nd.array(covariance, dtype=dtype))
rand_gen = None
var = MultivariateNormal.define_variable(shape=rv_shape, rand_gen=rand_gen, dtype=dtype).factor
variables = {var.mean.uuid: mean_mx, var.covariance.uuid: covariance_mx}
rv_samples_rt = var.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples)
assert array_has_samples(mx.nd, rv_samples_rt)
assert get_num_samples(mx.nd, rv_samples_rt) == num_samples
assert rv_samples_rt.dtype == dtype
from scipy.stats import multivariate_normal
if plot:
plot_bivariate(samples=rv_samples_rt, dist=multivariate_normal, mean=mean, cov=covariance)
# mean_est, scale_est = multivariate_normal.fit(rv_samples_rt.asnumpy().ravel())
mean_est = np.mean(rv_samples_rt.asnumpy(), axis=0)
cov_est = np.cov(rv_samples_rt.asnumpy(), rowvar=False)
mean_tol = 1e-1
covariance_tol = 1e-1
assert np.allclose(mean, mean_est, atol=mean_tol)
assert np.allclose(covariance, cov_est, covariance_tol)