def test_draw_samples_no_broadcast(self, dtype, mean, mean_is_samples, precision,
precision_is_samples, rv_shape, num_samples):
mean_mx = mx.nd.array(mean, dtype=dtype)
if not mean_is_samples:
mean_mx = add_sample_dimension(mx.nd, mean_mx)
precision_mx = mx.nd.array(precision, dtype=dtype)
if not precision_is_samples:
precision_mx = add_sample_dimension(mx.nd, precision_mx)
# precision = precision_mx.asnumpy()
# n_dim = 1 + len(rv.shape) if is_samples_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(precision)
# temp1 = np.matmul(lmat, rand_exp).sum(-1)
# rv_samples_np = mean + temp1
normal = MultivariateNormalMeanPrecision.define_variable(shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor
variables = {normal.mean.uuid: mean_mx, normal.precision.uuid: precision_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)
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_no_broadcast(self, dtype_dof, dtype,
degrees_of_freedom, scale,
scale_is_samples, rv_shape,
num_samples):
degrees_of_freedom_mx = mx.nd.array([degrees_of_freedom],
dtype=dtype_dof)
scale_mx = mx.nd.array(scale, dtype=dtype)
if not scale_is_samples:
scale_mx = add_sample_dimension(mx.nd, scale_mx)
rand = np.random.rand(num_samples, *rv_shape)
rand_gen = MockMXNetRandomGenerator(
mx.nd.array(rand.flatten(), dtype=dtype))
var = Wishart.define_variable(shape=rv_shape,
dtype=dtype,
rand_gen=rand_gen).factor
variables = {
var.degrees_of_freedom.uuid: degrees_of_freedom_mx,
var.scale.uuid: scale_mx
}
draw_samples_rt = var.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)
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_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_draw_samples_with_broadcast(self, dtype_dof, dtype, degrees_of_freedom, scale, scale_is_samples, rv_shape,
num_samples):
degrees_of_freedom_mx = mx.nd.array([degrees_of_freedom], dtype=dtype_dof)
scale_mx = mx.nd.array(scale, dtype=dtype)
if not scale_is_samples:
scale_mx = add_sample_dimension(mx.nd, scale_mx)
rand = np.random.rand(num_samples, *rv_shape)
rand_gen = MockMXNetRandomGenerator(mx.nd.array(rand.flatten(), dtype=dtype))
reps = 1000
mins = np.zeros(reps)
maxs = np.zeros(reps)
for i in range(reps):
rvs = wishart.rvs(df=degrees_of_freedom, scale=scale, size=num_samples)
mins[i] = rvs.min()
maxs[i] = rvs.max()
# rv_samples_np = wishart.rvs(df=degrees_of_freedom, scale=scale, size=num_samples)
var = Wishart.define_variable(shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor
variables = {var.degrees_of_freedom.uuid: degrees_of_freedom_mx, var.scale.uuid: scale_mx}
draw_samples_rt = var.draw_samples(F=mx.nd, variables=variables)
assert np.issubdtype(draw_samples_rt.dtype, dtype)
assert is_sampled_array(mx.nd, draw_samples_rt) == scale_is_samples
if scale_is_samples:
assert get_num_samples(mx.nd, draw_samples_rt) == num_samples, (get_num_samples(mx.nd, draw_samples_rt),
num_samples)
assert mins.min() < draw_samples_rt.asnumpy().min()
assert maxs.max() > draw_samples_rt.asnumpy().max()
def test_log_pdf(self, dtype_dof, dtype, degrees_of_freedom, random_state,
scale_is_samples, rv_is_samples, num_data_points, num_samples, broadcast):
# Create positive semi-definite matrices
rv = make_spd_matrices_4d(num_samples, num_data_points, degrees_of_freedom, random_state=random_state)
if broadcast:
scale = make_spd_matrix(n_dim=degrees_of_freedom, random_state=random_state)
else:
scale = make_spd_matrices_4d(num_samples, num_data_points, degrees_of_freedom, random_state=random_state)
degrees_of_freedom_mx = mx.nd.array([degrees_of_freedom], dtype=dtype_dof)
degrees_of_freedom = degrees_of_freedom_mx.asnumpy()[0] # ensures the correct dtype
scale_mx = mx.nd.array(scale, dtype=dtype)
if not scale_is_samples:
scale_mx = add_sample_dimension(mx.nd, scale_mx)
scale = scale_mx.asnumpy()
rv_mx = mx.nd.array(rv, dtype=dtype)
if not rv_is_samples:
rv_mx = add_sample_dimension(mx.nd, rv_mx)
rv = rv_mx.asnumpy()
is_samples_any = scale_is_samples or rv_is_samples
if broadcast:
scale_np = np.broadcast_to(scale, rv.shape)
else:
n_dim = 1 + len(rv.shape) if is_samples_any and not rv_is_samples else len(rv.shape)
scale_np = numpy_array_reshape(scale, is_samples_any, n_dim)
rv_np = numpy_array_reshape(rv, is_samples_any, degrees_of_freedom)
r = []
for s in range(num_samples):
a = []
for i in range(num_data_points):
a.append(wishart.logpdf(rv_np[s][i], df=degrees_of_freedom, scale=scale_np[s][i]))
r.append(a)
log_pdf_np = np.array(r)
var = Wishart.define_variable(shape=rv.shape[1:], dtype=dtype, rand_gen=None).factor
variables = {var.degrees_of_freedom.uuid: degrees_of_freedom_mx, var.scale.uuid: scale_mx,
var.random_variable.uuid: rv_mx}
log_pdf_rt = var.log_pdf(F=mx.nd, variables=variables)
assert np.issubdtype(log_pdf_rt.dtype, dtype)
assert is_sampled_array(mx.nd, log_pdf_rt) == is_samples_any
if is_samples_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_log_pdf(self, dtype, prob_true, prob_true_is_samples, rv,
rv_is_samples, num_samples):
rv_shape = rv.shape[1:] if rv_is_samples else rv.shape
n_dim = 1 + len(rv.shape) if not rv_is_samples else len(rv.shape)
prob_true_np = numpy_array_reshape(prob_true, prob_true_is_samples,
n_dim)
rv_np = numpy_array_reshape(rv, rv_is_samples, n_dim)
rv_full_shape = (num_samples, ) + rv_shape
rv_np = np.broadcast_to(rv_np, rv_full_shape)
log_pdf_np = bernoulli.logpmf(k=rv_np, p=prob_true_np)
var = Bernoulli.define_variable(0, shape=rv_shape, dtype=dtype).factor
prob_true_mx = mx.nd.array(prob_true, dtype=dtype)
if not prob_true_is_samples:
prob_true_mx = add_sample_dimension(mx.nd, prob_true_mx)
rv_mx = mx.nd.array(rv, dtype=dtype)
if not rv_is_samples:
rv_mx = add_sample_dimension(mx.nd, rv_mx)
variables = {
var.prob_true.uuid: prob_true_mx,
var.random_variable.uuid: rv_mx
}
log_pdf_rt = var.log_pdf(F=mx.nd, variables=variables)
assert np.issubdtype(log_pdf_rt.dtype, dtype)
assert get_num_samples(mx.nd, log_pdf_rt) == num_samples
assert np.allclose(log_pdf_np, log_pdf_rt.asnumpy())
def test_draw_samples(self, dtype, mean, mean_is_samples, precision,
precision_is_samples, rv_shape, num_samples):
n_dim = 1 + len(rv_shape)
mean_np = numpy_array_reshape(mean, mean_is_samples, n_dim)
precision_np = numpy_array_reshape(precision, precision_is_samples, n_dim)
rand = np.random.randn(num_samples, *rv_shape)
rv_samples_np = mean_np + rand * np.power(precision_np, -0.5)
rand_gen = MockMXNetRandomGenerator(mx.nd.array(rand.flatten(), dtype=dtype))
var = NormalMeanPrecision.define_variable(shape=rv_shape, dtype=dtype,
rand_gen=rand_gen).factor
mean_mx = mx.nd.array(mean, dtype=dtype)
if not mean_is_samples:
mean_mx = add_sample_dimension(mx.nd, mean_mx)
precision_mx = mx.nd.array(precision, dtype=dtype)
if not precision_is_samples:
precision_mx = add_sample_dimension(mx.nd, precision_mx)
variables = {var.mean.uuid: mean_mx, var.precision.uuid: precision_mx}
rv_samples_rt = var.draw_samples(
F=mx.nd, variables=variables, num_samples=num_samples)
assert np.issubdtype(rv_samples_rt.dtype, dtype)
assert array_has_samples(mx.nd, rv_samples_rt)
assert get_num_samples(mx.nd, rv_samples_rt) == num_samples
if np.issubdtype(dtype, np.float64):
rtol, atol = 1e-7, 1e-10
else:
rtol, atol = 1e-4, 1e-5
assert np.allclose(rv_samples_np, rv_samples_rt.asnumpy(), rtol=rtol, atol=atol)
def test_log_pdf(self, dtype, mean, mean_is_samples, precision, precision_is_samples,
rv, rv_is_samples, num_samples):
is_samples_any = any([mean_is_samples, precision_is_samples, rv_is_samples])
rv_shape = rv.shape[1:] if rv_is_samples else rv.shape
n_dim = 1 + len(rv.shape) if is_samples_any and not rv_is_samples else len(rv.shape)
mean_np = numpy_array_reshape(mean, mean_is_samples, n_dim)
precision_np = numpy_array_reshape(precision, precision_is_samples, n_dim)
rv_np = numpy_array_reshape(rv, rv_is_samples, n_dim)
log_pdf_np = norm.logpdf(rv_np, mean_np, np.power(precision_np, -0.5))
var = NormalMeanPrecision.define_variable(shape=rv_shape, dtype=dtype).factor
mean_mx = mx.nd.array(mean, dtype=dtype)
if not mean_is_samples:
mean_mx = add_sample_dimension(mx.nd, mean_mx)
precision_mx = mx.nd.array(precision, dtype=dtype)
if not precision_is_samples:
precision_mx = add_sample_dimension(mx.nd, precision_mx)
rv_mx = mx.nd.array(rv, dtype=dtype)
if not rv_is_samples:
rv_mx = add_sample_dimension(mx.nd, rv_mx)
variables = {var.mean.uuid: mean_mx, var.precision.uuid: precision_mx, var.random_variable.uuid: rv_mx}
log_pdf_rt = var.log_pdf(F=mx.nd, variables=variables)
assert np.issubdtype(log_pdf_rt.dtype, dtype)
assert array_has_samples(mx.nd, log_pdf_rt) == is_samples_any
if is_samples_any:
assert get_num_samples(mx.nd, log_pdf_rt) == num_samples
if np.issubdtype(dtype, np.float64):
rtol, atol = 1e-7, 1e-10
else:
rtol, atol = 1e-4, 1e-5
assert np.allclose(log_pdf_np, log_pdf_rt.asnumpy(), rtol=rtol, atol=atol)
def test_draw_samples(self, dtype, mean, mean_isSamples, var,
var_isSamples, rv_shape, num_samples):
n_dim = 1 + len(rv_shape)
mean_np = numpy_array_reshape(mean, mean_isSamples, n_dim)
var_np = numpy_array_reshape(var, var_isSamples, n_dim)
rand = np.random.randn(num_samples, *rv_shape)
rv_samples_np = mean_np + rand * np.sqrt(var_np)
rand_gen = MockMXNetRandomGenerator(mx.nd.array(rand.flatten(), dtype=dtype))
normal = Normal.define_variable(shape=rv_shape, dtype=dtype,
rand_gen=rand_gen).factor
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)
variables = {normal.mean.uuid: mean_mx, normal.variance.uuid: var_mx}
rv_samples_rt = normal.draw_samples(
F=mx.nd, variables=variables, num_samples=num_samples)
assert np.issubdtype(rv_samples_rt.dtype, dtype)
assert array_has_samples(mx.nd, rv_samples_rt)
assert get_num_samples(mx.nd, rv_samples_rt) == num_samples
if np.issubdtype(dtype, np.float64):
rtol, atol = 1e-7, 1e-10
else:
rtol, atol = 1e-4, 1e-5
assert np.allclose(rv_samples_np, rv_samples_rt.asnumpy(), rtol=rtol, atol=atol)
def test_log_pdf(self, dtype, mean, mean_isSamples, var, var_isSamples,
rv, rv_isSamples, num_samples):
from scipy.stats import norm
isSamples_any = any([mean_isSamples, var_isSamples, rv_isSamples])
rv_shape = rv.shape[1:] if rv_isSamples else rv.shape
n_dim = 1 + len(rv.shape) if isSamples_any and not rv_isSamples else len(rv.shape)
mean_np = numpy_array_reshape(mean, mean_isSamples, n_dim)
var_np = numpy_array_reshape(var, var_isSamples, n_dim)
rv_np = numpy_array_reshape(rv, rv_isSamples, n_dim)
log_pdf_np = norm.logpdf(rv_np, mean_np, np.sqrt(var_np))
normal = Normal.define_variable(shape=rv_shape, dtype=dtype).factor
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)
rv_mx = mx.nd.array(rv, dtype=dtype)
if not rv_isSamples:
rv_mx = add_sample_dimension(mx.nd, rv_mx)
variables = {normal.mean.uuid: mean_mx, normal.variance.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 array_has_samples(mx.nd, log_pdf_rt) == isSamples_any
if isSamples_any:
assert get_num_samples(mx.nd, log_pdf_rt) == num_samples
if np.issubdtype(dtype, np.float64):
rtol, atol = 1e-7, 1e-10
else:
rtol, atol = 1e-4, 1e-5
assert np.allclose(log_pdf_np, log_pdf_rt.asnumpy(), rtol=rtol, atol=atol)
def test_draw_samples_non_mock(self, plot=False):
# Also make sure the non-mock sampler works
dtype = np.float32
num_samples = 100000
a = np.array([2])
b = np.array([5])
rv_shape = (1, )
a_mx = add_sample_dimension(mx.nd, mx.nd.array(a, dtype=dtype))
b_mx = add_sample_dimension(mx.nd, mx.nd.array(b, dtype=dtype))
rand_gen = None
var = Beta.define_variable(shape=rv_shape,
rand_gen=rand_gen,
dtype=dtype).factor
variables = {var.alpha.uuid: a_mx, var.beta.uuid: b_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
if plot:
plot_univariate(samples=rv_samples_rt, dist=beta, a=a[0], b=b[0])
a_est, b_est, _, _ = beta.fit(rv_samples_rt.asnumpy().ravel())
a_tol = 0.2
b_tol = 0.2
assert np.abs(a[0] - a_est) < a_tol
assert np.abs(b[0] - b_est) < b_tol
def test_draw_samples(self, dtype, low, low_is_samples, high,
high_is_samples, rv_shape, num_samples):
n_dim = 1 + len(rv_shape)
low_np = numpy_array_reshape(low, low_is_samples, n_dim)
high_np = numpy_array_reshape(high, high_is_samples, n_dim)
rv_samples_np = np.random.uniform(low=low_np, high=high_np, size=(num_samples,) + rv_shape)
rand_gen = MockMXNetRandomGenerator(mx.nd.array(rv_samples_np.flatten(), dtype=dtype))
var = Uniform.define_variable(shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor
low_mx = mx.nd.array(low, dtype=dtype)
if not low_is_samples:
low_mx = add_sample_dimension(mx.nd, low_mx)
high_mx = mx.nd.array(high, dtype=dtype)
if not high_is_samples:
high_mx = add_sample_dimension(mx.nd, high_mx)
variables = {var.low.uuid: low_mx, var.high.uuid: high_mx}
rv_samples_rt = var.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples)
assert np.issubdtype(rv_samples_rt.dtype, dtype)
assert array_has_samples(mx.nd, rv_samples_rt)
assert get_num_samples(mx.nd, rv_samples_rt) == num_samples
if np.issubdtype(dtype, np.float64):
rtol, atol = 1e-7, 1e-10
else:
rtol, atol = 1e-4, 1e-5
assert np.allclose(rv_samples_np, rv_samples_rt.asnumpy(), rtol=rtol, atol=atol)
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])
variance = np.array([2])
rv_shape = (1,)
mean_mx = add_sample_dimension(mx.nd, mx.nd.array(mean, dtype=dtype))
variance_mx = add_sample_dimension(mx.nd, mx.nd.array(variance, dtype=dtype))
rand_gen = None
var = Normal.define_variable(shape=rv_shape, rand_gen=rand_gen, dtype=dtype).factor
variables = {var.mean.uuid: mean_mx, var.variance.uuid: variance_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 norm
if plot:
plot_univariate(samples=rv_samples_rt, dist=norm, loc=mean[0], scale=np.sqrt(variance[0]))
mean_est, scale_est = norm.fit(rv_samples_rt.asnumpy().ravel())
mean_tol = 1e-1
variance_tol = 1e-1
assert np.abs(mean[0] - mean_est) < mean_tol
assert np.abs(variance[0] - scale_est ** 2) < variance_tol
def test_draw_samples_non_mock(self, plot=False):
# Also make sure the non-mock sampler works
dtype = np.float32
num_samples = 1000
low = np.array([0.5])
high = np.array([2])
rv_shape = (1,)
low_mx = mx.nd.array(low, dtype=dtype)
high_mx = mx.nd.array(high, dtype=dtype)
rand_gen = None
var = Uniform.define_variable(shape=rv_shape, rand_gen=rand_gen, dtype=dtype).factor
variables = {var.low.uuid: low_mx, var.high.uuid: high_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
if plot:
plot_univariate(samples=rv_samples_rt, dist=uniform, loc=low[0], scale=high[0] - low[0], buffer=1)
location_est, scale_est = uniform.fit(rv_samples_rt.asnumpy().ravel())
location_tol = 1e-2
scale_tol = 1e-2
assert np.abs(low[0] - location_est) < location_tol
assert np.abs(high[0] - scale_est - location_est) < scale_tol
def test_draw_samples(self):
np.random.seed(0)
samples_1_np = np.random.randn(5)
samples_1 = mx.nd.array(samples_1_np)
samples_2_np = np.random.randn(50)
samples_2 = mx.nd.array(samples_2_np)
m = Model()
v = Variable(shape=(1,))
m.v2 = Normal.define_variable(mean=v, variance=mx.nd.array([1]), rand_gen=MockMXNetRandomGenerator(samples_1))
m.v3 = Normal.define_variable(mean=m.v2, variance=mx.nd.array([0.1]), shape=(10,), rand_gen=MockMXNetRandomGenerator(samples_2))
np.random.seed(0)
v_np =np.random.rand(1)
v_mx = mx.nd.array(v_np)
v_rt = add_sample_dimension(mx.nd, v_mx)
variance = m.v2.factor.variance
variance2 = m.v3.factor.variance
variance_rt = add_sample_dimension(mx.nd, variance.constant)
variance2_rt = add_sample_dimension(mx.nd, variance2.constant)
samples = m.draw_samples(F=mx.nd, num_samples=5, targets=[m.v3.uuid],
variables={v.uuid: v_rt, variance.uuid: variance_rt, variance2.uuid: variance2_rt})[0]
samples_np = v_np + samples_1_np[:, None] + np.sqrt(0.1)*samples_2_np.reshape(5,10)
assert array_has_samples(mx.nd, samples) and get_num_samples(mx.nd, samples)==5
assert np.allclose(samples.asnumpy(), samples_np)
def test_log_pdf(self, dtype, low, low_is_samples, high, high_is_samples, rv, rv_is_samples,
num_samples):
is_samples_any = any([low_is_samples, high_is_samples, rv_is_samples])
rv_shape = rv.shape[1:] if rv_is_samples else rv.shape
n_dim = 1 + len(rv.shape) if is_samples_any and not rv_is_samples else len(rv.shape)
low_np = numpy_array_reshape(low, low_is_samples, n_dim)
high_np = numpy_array_reshape(high, high_is_samples, n_dim)
scale_np = high_np - low_np
rv_np = numpy_array_reshape(rv, rv_is_samples, n_dim)
# Note uniform.logpdf takes loc and scale, where loc=a and scale=b-a
log_pdf_np = uniform.logpdf(rv_np, low_np, scale_np)
var = Uniform.define_variable(shape=rv_shape, dtype=dtype).factor
low_mx = mx.nd.array(low, dtype=dtype)
if not low_is_samples:
low_mx = add_sample_dimension(mx.nd, low_mx)
high_mx = mx.nd.array(high, dtype=dtype)
if not high_is_samples:
high_mx = add_sample_dimension(mx.nd, high_mx)
rv_mx = mx.nd.array(rv, dtype=dtype)
if not rv_is_samples:
rv_mx = add_sample_dimension(mx.nd, rv_mx)
variables = {var.low.uuid: low_mx, var.high.uuid: high_mx, var.random_variable.uuid: rv_mx}
log_pdf_rt = var.log_pdf(F=mx.nd, variables=variables)
assert np.issubdtype(log_pdf_rt.dtype, dtype)
assert array_has_samples(mx.nd, log_pdf_rt) == is_samples_any
if is_samples_any:
assert get_num_samples(mx.nd, log_pdf_rt) == num_samples
if np.issubdtype(dtype, np.float64):
rtol, atol = 1e-7, 1e-10
else:
rtol, atol = 1e-4, 1e-5
assert np.allclose(log_pdf_np, log_pdf_rt.asnumpy(), rtol=rtol, atol=atol)
def test_draw_samples(self, dtype, log_prob, log_prob_isSamples, rv_shape,
num_samples, one_hot_encoding, normalization):
n_dim = 1 + len(rv_shape)
log_prob_np = numpy_array_reshape(log_prob, log_prob_isSamples, n_dim)
rv_full_shape = (num_samples, ) + rv_shape
log_prob_np = np.broadcast_to(log_prob_np, rv_full_shape[:-1] + (3, ))
rand_np = np.random.randint(0, 3, size=rv_full_shape[:-1])
rand_gen = MockMXNetRandomGenerator(
mx.nd.array(rand_np.flatten(), dtype=dtype))
if one_hot_encoding:
rand_np = np.identity(3)[rand_np].reshape(*rv_full_shape)
else:
rand_np = np.expand_dims(rand_np, axis=-1)
rv_samples_np = rand_np
cat = Categorical.define_variable(0,
num_classes=3,
one_hot_encoding=one_hot_encoding,
normalization=normalization,
shape=rv_shape,
rand_gen=rand_gen,
dtype=dtype).factor
log_prob_mx = mx.nd.array(log_prob, dtype=dtype)
if not log_prob_isSamples:
log_prob_mx = add_sample_dimension(mx.nd, log_prob_mx)
variables = {cat.log_prob.uuid: log_prob_mx}
rv_samples_rt = cat.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 np.allclose(rv_samples_np, rv_samples_rt.asnumpy())
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)
def test_log_pdf_with_broadcast(self, dtype, mean, mean_is_samples, precision, precision_is_samples,
rv, rv_is_samples, num_samples):
mean_mx = mx.nd.array(mean, dtype=dtype)
if not mean_is_samples:
mean_mx = add_sample_dimension(mx.nd, mean_mx)
mean = mean_mx.asnumpy()
precision_mx = mx.nd.array(precision, dtype=dtype)
if not precision_is_samples:
precision_mx = add_sample_dimension(mx.nd, precision_mx)
precision = precision_mx.asnumpy()
rv_mx = mx.nd.array(rv, dtype=dtype)
if not rv_is_samples:
rv_mx = add_sample_dimension(mx.nd, rv_mx)
rv = rv_mx.asnumpy()
is_samples_any = any([mean_is_samples, precision_is_samples, rv_is_samples])
rv_shape = rv.shape[1:]
n_dim = 1 + len(rv.shape) if is_samples_any and not rv_is_samples else len(rv.shape)
mean_np = np.broadcast_to(mean, (5, 3, 2))
precision_np = np.broadcast_to(precision, (5, 3, 2, 2))
rv_np = numpy_array_reshape(rv, is_samples_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], np.linalg.inv(precision_np[s][i])))
r.append(a)
log_pdf_np = np.array(r)
normal = MultivariateNormalMeanPrecision.define_variable(shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor
variables = {
normal.mean.uuid: mean_mx, normal.precision.uuid: precision_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 array_has_samples(mx.nd, log_pdf_rt) == is_samples_any
if is_samples_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(self, dtype, prob_true, prob_true_is_samples,
rv_shape, num_samples):
rv_full_shape = (num_samples, ) + rv_shape
rand_np = np.random.normal(size=rv_full_shape) > 0
rand_gen = MockMXNetRandomGenerator(
mx.nd.array(rand_np.flatten(), dtype=dtype))
rv_samples_np = rand_np
var = Bernoulli.define_variable(0,
shape=rv_shape,
rand_gen=rand_gen,
dtype=dtype).factor
prob_true_mx = mx.nd.array(prob_true, dtype=dtype)
if not prob_true_is_samples:
prob_true_mx = add_sample_dimension(mx.nd, prob_true_mx)
variables = {var.prob_true.uuid: prob_true_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 np.array_equal(rv_samples_np,
rv_samples_rt.asnumpy().astype(bool))
# Also make sure the non-mock sampler works
rand_gen = None
var = Bernoulli.define_variable(0,
shape=rv_shape,
rand_gen=rand_gen,
dtype=dtype).factor
prob_true_mx = mx.nd.array(prob_true, dtype=dtype)
if not prob_true_is_samples:
prob_true_mx = add_sample_dimension(mx.nd, prob_true_mx)
variables = {var.prob_true.uuid: prob_true_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
def test_clone_gp(self, dtype, X, X_isSamples, rbf_lengthscale,
rbf_lengthscale_isSamples, rbf_variance,
rbf_variance_isSamples, rv, rv_isSamples, num_samples):
X_mx = prepare_mxnet_array(X, X_isSamples, dtype)
rbf_lengthscale_mx = prepare_mxnet_array(rbf_lengthscale,
rbf_lengthscale_isSamples,
dtype)
rbf_variance_mx = prepare_mxnet_array(rbf_variance,
rbf_variance_isSamples, dtype)
rv_mx = prepare_mxnet_array(rv, rv_isSamples, dtype)
rv_shape = rv.shape[1:] if rv_isSamples else rv.shape
rbf = RBF(2, True, 1., 1., 'rbf', None, dtype)
m = Model()
m.X_var = Variable(shape=(5, 2))
m.Y = GaussianProcess.define_variable(X=m.X_var,
kernel=rbf,
shape=rv_shape,
dtype=dtype)
gp = m.clone().Y.factor
variables = {
gp.X.uuid: X_mx,
gp.rbf_lengthscale.uuid: rbf_lengthscale_mx,
gp.rbf_variance.uuid: rbf_variance_mx,
gp.random_variable.uuid: rv_mx
}
log_pdf_rt = gp.log_pdf(F=mx.nd, variables=variables).asnumpy()
log_pdf_np = []
for i in range(num_samples):
X_i = X[i] if X_isSamples else X
lengthscale_i = rbf_lengthscale[
i] if rbf_lengthscale_isSamples else rbf_lengthscale
variance_i = rbf_variance[
i] if rbf_variance_isSamples else rbf_variance
rv_i = rv[i] if rv_isSamples else rv
rbf_np = GPy.kern.RBF(input_dim=2, ARD=True)
rbf_np.lengthscale = lengthscale_i
rbf_np.variance = variance_i
K_np = rbf_np.K(X_i)
log_pdf_np.append(
multivariate_normal.logpdf(rv_i[:, 0], mean=None, cov=K_np))
log_pdf_np = np.array(log_pdf_np)
isSamples_any = any([
X_isSamples, rbf_lengthscale_isSamples, rbf_variance_isSamples,
rv_isSamples
])
assert np.issubdtype(log_pdf_rt.dtype, dtype)
assert array_has_samples(mx.nd, log_pdf_rt) == isSamples_any
if isSamples_any:
assert get_num_samples(mx.nd, log_pdf_rt) == num_samples
assert np.allclose(log_pdf_np, log_pdf_rt)
def test_log_pdf_with_broadcast(self, dtype, a, a_is_samples, rv, rv_is_samples, num_samples):
# Add sample dimension if varaible is not samples
a_mx = mx.nd.array(a, dtype=dtype)
if not a_is_samples:
a_mx = add_sample_dimension(mx.nd, a_mx)
a = a_mx.asnumpy()
rv_mx = mx.nd.array(rv, dtype=dtype)
if not rv_is_samples:
rv_mx = add_sample_dimension(mx.nd, rv_mx)
rv = rv_mx.asnumpy()
is_samples_any = a_is_samples or rv_is_samples
rv_shape = rv.shape[1:]
n_dim = 1 + len(rv.shape) if is_samples_any and not rv_is_samples else len(rv.shape)
a_np = np.broadcast_to(a, (num_samples, 3, 2))
rv_np = numpy_array_reshape(rv, is_samples_any, n_dim)
# Initialize rand_gen
rand = np.random.rand(num_samples, *rv_shape)
rand_gen = MockMXNetRandomGenerator(mx.nd.array(rand.flatten(), dtype=dtype))
# Calculate correct Dirichlet logpdf
r = []
for s in range(len(rv_np)):
a = []
for i in range(len(rv_np[s])):
a.append(scipy_dirichlet.logpdf(rv_np[s][i]/sum(rv_np[s][i]), a_np[s][i]))
r.append(a)
log_pdf_np = np.array(r)
dirichlet = Dirichlet.define_variable(alpha=Variable(), shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor
variables = {dirichlet.alpha.uuid: a_mx, dirichlet.random_variable.uuid: rv_mx}
log_pdf_rt = dirichlet.log_pdf(F=mx.nd, variables=variables)
assert np.issubdtype(log_pdf_rt.dtype, dtype)
assert array_has_samples(mx.nd, log_pdf_rt) == is_samples_any
if is_samples_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_mean_variance(self, dtype, mean, mean_isSamples,
variance, variance_isSamples, rv_shape,
num_samples):
n_dim = 1 + len(rv_shape)
out_shape = (num_samples, ) + rv_shape
mean_np = mx.nd.array(np.broadcast_to(numpy_array_reshape(
mean, mean_isSamples, n_dim),
shape=out_shape),
dtype=dtype)
variance_np = mx.nd.array(np.broadcast_to(numpy_array_reshape(
variance, variance_isSamples, n_dim),
shape=out_shape),
dtype=dtype)
gamma = GammaMeanVariance.define_variable(shape=rv_shape,
dtype=dtype).factor
mean_mx = mx.nd.array(mean, dtype=dtype)
if not mean_isSamples:
mean_mx = add_sample_dimension(mx.nd, mean_mx)
variance_mx = mx.nd.array(variance, dtype=dtype)
if not variance_isSamples:
variance_mx = add_sample_dimension(mx.nd, variance_mx)
variables = {
gamma.mean.uuid: mean_mx,
gamma.variance.uuid: variance_mx
}
mx.random.seed(0)
rv_samples_rt = gamma.draw_samples(F=mx.nd,
variables=variables,
num_samples=num_samples)
mx.random.seed(0)
beta_np = mean_np / variance_np
alpha_np = mean_np * beta_np
rv_samples_mx = mx.nd.random.gamma(alpha=alpha_np,
beta=beta_np,
dtype=dtype)
assert np.issubdtype(rv_samples_rt.dtype, dtype)
assert is_sampled_array(mx.nd, rv_samples_rt)
assert get_num_samples(mx.nd, rv_samples_rt) == num_samples
if np.issubdtype(dtype, np.float64):
rtol, atol = 1e-7, 1e-10
else:
rtol, atol = 1e-4, 1e-5
assert np.allclose(rv_samples_mx.asnumpy(),
rv_samples_rt.asnumpy(),
rtol=rtol,
atol=atol)
def test_log_pdf_mean_variance(self, dtype, mean, mean_isSamples, variance,
variance_isSamples, rv, rv_isSamples,
num_samples):
import scipy as sp
isSamples_any = any([mean_isSamples, variance_isSamples, rv_isSamples])
rv_shape = rv.shape[1:] if rv_isSamples else rv.shape
n_dim = 1 + len(
rv.shape) if isSamples_any and not rv_isSamples else len(rv.shape)
mean_np = numpy_array_reshape(mean, mean_isSamples, n_dim)
variance_np = numpy_array_reshape(variance, variance_isSamples, n_dim)
rv_np = numpy_array_reshape(rv, rv_isSamples, n_dim)
beta_np = mean_np / variance_np
alpha_np = mean_np * beta_np
log_pdf_np = sp.stats.gamma.logpdf(rv_np,
a=alpha_np,
loc=0,
scale=1. / beta_np)
mean_mx = mx.nd.array(mean, dtype=dtype)
if not mean_isSamples:
mean_mx = add_sample_dimension(mx.nd, mean_mx)
variance_mx = mx.nd.array(variance, dtype=dtype)
if not variance_isSamples:
variance_mx = add_sample_dimension(mx.nd, variance_mx)
rv_mx = mx.nd.array(rv, dtype=dtype)
if not rv_isSamples:
rv_mx = add_sample_dimension(mx.nd, rv_mx)
gamma = GammaMeanVariance.define_variable(mean=mean_mx,
variance=variance_mx,
shape=rv_shape,
dtype=dtype).factor
variables = {
gamma.mean.uuid: mean_mx,
gamma.variance.uuid: variance_mx,
gamma.random_variable.uuid: rv_mx
}
log_pdf_rt = gamma.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
if np.issubdtype(dtype, np.float64):
rtol, atol = 1e-7, 1e-10
else:
rtol, atol = 1e-4, 1e-5
assert np.allclose(log_pdf_np,
log_pdf_rt.asnumpy(),
rtol=rtol,
atol=atol)
def test_draw_samples(self, dtype, X, X_isSamples, rbf_lengthscale,
rbf_lengthscale_isSamples, rbf_variance,
rbf_variance_isSamples, rv_shape, num_samples):
X_mx = prepare_mxnet_array(X, X_isSamples, dtype)
rbf_lengthscale_mx = prepare_mxnet_array(rbf_lengthscale,
rbf_lengthscale_isSamples,
dtype)
rbf_variance_mx = prepare_mxnet_array(rbf_variance,
rbf_variance_isSamples, dtype)
rand = np.random.randn(num_samples, *rv_shape)
rand_gen = MockMXNetRandomGenerator(
mx.nd.array(rand.flatten(), dtype=dtype))
rbf = RBF(2, True, 1., 1., 'rbf', None, dtype)
X_var = Variable(shape=(5, 2))
gp = GaussianProcess.define_variable(X=X_var,
kernel=rbf,
shape=rv_shape,
dtype=dtype,
rand_gen=rand_gen).factor
variables = {
gp.X.uuid: X_mx,
gp.rbf_lengthscale.uuid: rbf_lengthscale_mx,
gp.rbf_variance.uuid: rbf_variance_mx
}
samples_rt = gp.draw_samples(F=mx.nd,
variables=variables,
num_samples=num_samples).asnumpy()
samples_np = []
for i in range(num_samples):
X_i = X[i] if X_isSamples else X
lengthscale_i = rbf_lengthscale[
i] if rbf_lengthscale_isSamples else rbf_lengthscale
variance_i = rbf_variance[
i] if rbf_variance_isSamples else rbf_variance
rand_i = rand[i]
rbf_np = GPy.kern.RBF(input_dim=2, ARD=True)
rbf_np.lengthscale = lengthscale_i
rbf_np.variance = variance_i
K_np = rbf_np.K(X_i)
L_np = np.linalg.cholesky(K_np)
sample_np = L_np.dot(rand_i)
samples_np.append(sample_np)
samples_np = np.array(samples_np)
assert np.issubdtype(samples_rt.dtype, dtype)
assert get_num_samples(mx.nd, samples_rt) == num_samples
assert np.allclose(samples_np, samples_rt)
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()
# var = (var[:,:,:,None]*var[:,None,:,:]).sum(-2)+np.eye(2)
# var_mx = mx.nd.array(var, dtype=dtype)
isSamples_any = any([mean_isSamples, var_isSamples])
# n_dim = 1 + len(rv.shape) if isSamples_any else len(rv.shape)
# mean_np = numpy_array_reshape(mean, mean_isSamples, n_dim)
# var_np = numpy_array_reshape(var, var_isSamples, n_dim)
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)
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_draw_samples_no_broadcast(self, dtype_dof, dtype, degrees_of_freedom, scale,
scale_is_samples, rv_shape, num_samples):
degrees_of_freedom_mx = mx.nd.array([degrees_of_freedom], dtype=dtype_dof)
scale_mx = mx.nd.array(scale, dtype=dtype)
if not scale_is_samples:
scale_mx = add_sample_dimension(mx.nd, scale_mx)
# 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
var = Wishart.define_variable(shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor
variables = {var.degrees_of_freedom.uuid: degrees_of_freedom_mx, var.scale.uuid: scale_mx}
draw_samples_rt = var.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)
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(self, dtype, a_shape, a_is_samples, b_shape,
b_is_samples, rv_shape, num_samples):
# Note: Tests above have been commented as they are very slow to run.
# Note: Moved random number generation to here as the seed wasn't set if used above
a = np.random.uniform(0.5, 2, size=a_shape)
b = np.random.uniform(0.5, 2, size=b_shape)
n_dim = 1 + len(rv_shape)
a_np = numpy_array_reshape(a, a_is_samples, n_dim)
b_np = numpy_array_reshape(b, b_is_samples, n_dim)
rv_samples_np = np.random.beta(a_np,
b_np,
size=(num_samples, ) + rv_shape)
var = Beta.define_variable(shape=rv_shape, dtype=dtype,
rand_gen=None).factor
a_mx = mx.nd.array(a, dtype=dtype)
if not a_is_samples:
a_mx = add_sample_dimension(mx.nd, a_mx)
b_mx = mx.nd.array(b, dtype=dtype)
if not b_is_samples:
b_mx = add_sample_dimension(mx.nd, b_mx)
variables = {var.a.uuid: a_mx, var.b.uuid: b_mx}
rv_samples_rt = var.draw_samples(F=mx.nd,
variables=variables,
num_samples=num_samples)
assert np.issubdtype(rv_samples_rt.dtype, dtype)
assert is_sampled_array(mx.nd, rv_samples_rt)
assert get_num_samples(mx.nd, rv_samples_rt) == num_samples
rtol, atol = 1e-1, 1e-1
from itertools import product
fits_np = [
beta.fit(rv_samples_np[:, i, j])[0:2]
for i, j in (product(*map(range, rv_shape)))
]
fits_rt = [
beta.fit(rv_samples_rt.asnumpy()[:, i, j])[0:2]
for i, j in (product(*map(range, rv_shape)))
]
assert np.allclose(fits_np, fits_rt, rtol=rtol, atol=atol)
def test_log_pdf(self, dtype, log_prob, log_prob_isSamples, rv,
rv_isSamples, num_samples, one_hot_encoding,
normalization):
rv_shape = rv.shape[1:] if rv_isSamples else rv.shape
n_dim = 1 + len(rv.shape) if not rv_isSamples else len(rv.shape)
log_prob_np = numpy_array_reshape(log_prob, log_prob_isSamples, n_dim)
rv_np = numpy_array_reshape(rv, rv_isSamples, n_dim)
rv_full_shape = (num_samples, ) + rv_shape
rv_np = np.broadcast_to(rv_np, rv_full_shape)
log_prob_np = np.broadcast_to(log_prob_np, rv_full_shape[:-1] + (3, ))
if normalization:
log_pdf_np = np.log(
np.exp(log_prob_np) /
np.exp(log_prob_np).sum(-1, keepdims=True)).reshape(-1, 3)
else:
log_pdf_np = log_prob_np.reshape(-1, 3)
if one_hot_encoding:
log_pdf_np = (rv_np.reshape(-1, 3) * log_pdf_np).sum(-1).reshape(
rv_np.shape[:-1])
else:
bool_idx = np.arange(3)[None, :] == rv_np.reshape(-1, 1)
log_pdf_np = log_pdf_np[bool_idx].reshape(rv_np.shape[:-1])
cat = Categorical.define_variable(0,
num_classes=3,
one_hot_encoding=one_hot_encoding,
normalization=normalization,
shape=rv_shape,
dtype=dtype).factor
log_prob_mx = mx.nd.array(log_prob, dtype=dtype)
if not log_prob_isSamples:
log_prob_mx = add_sample_dimension(mx.nd, log_prob_mx)
rv_mx = mx.nd.array(rv, dtype=dtype)
if not rv_isSamples:
rv_mx = add_sample_dimension(mx.nd, rv_mx)
variables = {
cat.log_prob.uuid: log_prob_mx,
cat.random_variable.uuid: rv_mx
}
log_pdf_rt = cat.log_pdf(F=mx.nd, variables=variables)
assert np.issubdtype(log_pdf_rt.dtype, dtype)
assert get_num_samples(mx.nd, log_pdf_rt) == num_samples
assert np.allclose(log_pdf_np, log_pdf_rt.asnumpy())