def test_CSGraph__init_transition_probs(g):
# Gives the right values upon class instanciation
set(g._edge_aliases.keys()) == {(0, 1), (0, 2), (0, 3), (1, 0), (1, 3),
(1, 5), (2, 0), (2, 3), (2, 5), (3, 0),
(3, 1), (3, 2), (3, 4), (4, 3), (4, 5),
(5, 1), (5, 2), (5, 4)}
probs20 = np.array([1 / 4, 2, 1])
assert_array_equal(g._edge_aliases[(2, 0)],
utils.alias_setup(probs20 / probs20.sum()))
probs31 = np.array([1, 2, 1 / 4])
assert_array_equal(g._edge_aliases[(3, 1)],
utils.alias_setup(probs31 / probs31.sum()))
def test_alias_draw():
# Works with good arguments
probs = np.arange(5)
probs = probs / probs.sum()
choices, weights = utils.alias_setup(probs)
assert utils.alias_draw(choices, weights) in range(1, 5)
# Gives the right statistics
n_draws = 1000
draws = np.array(
[utils.alias_draw(choices, weights) for _ in range(n_draws)])
counts = [(draws == i).sum() for i in range(5)]
assert counts[0] == 0
assert np.sum(counts[1:]) == n_draws
assert_allclose(counts[1:], n_draws * probs[1:], atol=50)
# Fails with empty or None inputs
with pytest.raises(numba.errors.TypingError):
utils.alias_draw(None, weights)
with pytest.raises(numba.errors.TypingError):
utils.alias_draw(choices, None)
with pytest.raises(numba.errors.TypingError):
utils.alias_draw(None, None)
with pytest.raises(AssertionError):
utils.alias_draw(choices[:0], weights)
with pytest.raises(AssertionError):
utils.alias_draw(choices, weights[:0])
with pytest.raises(AssertionError):
utils.alias_draw(choices[:0], weights[:0])
def test_alias_setup():
# Fails if `probs` is not an ndarray or is empty
with pytest.raises(numba.errors.TypingError):
utils.alias_setup([.5, .5])
with pytest.raises(AssertionError):
utils.alias_setup(np.array([]))
# Works with good arguments
probs = np.ones(4)
probs = probs / probs.sum()
choices, weights = utils.alias_setup(probs)
assert_array_equal(choices, [0, 0, 0, 0])
assert_array_equal(weights, [1, 1, 1, 1])
probs = np.array([1, 1, 2, 2])
probs = probs / probs.sum()
choices, weights = utils.alias_setup(probs)
assert_array_equal(choices, [3, 3, 0, 2])
assert_allclose(weights, [2 / 3, 2 / 3, 1, 2 / 3], rtol=1e-6)
probs = np.arange(5)
probs = probs / probs.sum()
choices, weights = utils.alias_setup(probs)
assert_array_equal(choices, [4, 4, 0, 0, 3])
assert_array_equal(weights, [0, .5, 1, 1, .5])
def stbm(n_nodes=100, n_clusters=40, # Network parameters
n_topics=5, n_documents=30, n_slots=140, vocabulary_size=50, # Language parameters
ρ=None, π=None, θ=None):
"""Simulate a network + document scenario with STBM."""
β = np.random.random(size=(n_topics, vocabulary_size))
β /= np.expand_dims(np.sum(β, axis=1), -1)
if ρ is None:
ρ = np.random.random(n_clusters)
ρ /= np.sum(ρ)
if π is None:
pre_π = np.random.random((n_clusters, n_clusters))
π = np.triu(pre_π, k=1) + np.triu(pre_π, k=0).T
# θ cannot be None, but we can't enforce that with Numba, it'll just bail further down
# Order Latent Variables
Y = np.random.multinomial(1, pvals=ρ, size=n_nodes).astype(np.float_)
A_probs = Y @ π @ Y.T
A = np.zeros_like(A_probs)
for i in numba.prange(len(A_probs.flat)):
A.flat[i] = np.random.binomial(1, A_probs.flat[i])
Z = np.nan * np.zeros((n_nodes, n_nodes, n_documents, n_slots))
W = np.nan * np.zeros((n_nodes, n_nodes, n_documents, n_slots))
sources_A, destinations_A = np.where(A)
for i, j in zip(sources_A, destinations_A):
qs, rs = np.where(Y[i])[0], np.where(Y[j])[0]
assert len(qs) == len(rs) == 1
q, r = qs[0], rs[0]
topics_choice_aliases = utils.alias_setup(θ[q, r])
for k in numba.prange(n_documents * n_slots):
Z[i, j].flat[k] = utils.alias_draw(*topics_choice_aliases)
for k in numba.prange(n_topics):
k_idx = np.where(Z == k)
vocabulary_choice_aliases = utils.alias_setup(β[k])
for idx in zip(*k_idx):
W[idx] = utils.alias_draw(*vocabulary_choice_aliases)
return (Y, A, Z, W)
def test_CSGraph__get_alias_edge(g):
# Works well with good parameters
probs01 = np.array([2, 1, 1 / 4])
assert_alias_equal(g._get_alias_edge(0, 1),
utils.alias_setup(probs01 / probs01.sum()))
# Fails if `dst` has no neighbours (which should never happen,
# since we always arrive through a link)
g.adj = sparse.csr_matrix(
np.array([[0, 1, 1, 1, 0, 0], [1, 0, 0, 1, 0, 0], [1, 0, 0, 1, 0, 0],
[1, 1, 1, 0, 1, 0], [0, 0, 0, 1, 0, 0], [0, 0, 0, 0, 0, 0]]))
with pytest.raises(AssertionError):
g._get_alias_edge(1, 5)
def _get_alias_edge(src, dst, p, q, adj_data, adj_indices, adj_indptr,
adj_shape):
src_neigbours = _neighbours(src, adj_data, adj_indices, adj_indptr,
adj_shape)
dst_neigbours = _neighbours(dst, adj_data, adj_indices, adj_indptr,
adj_shape)
n_dst_neigbours = len(dst_neigbours)
assert n_dst_neigbours > 0
uprobs = np.ones(n_dst_neigbours) / q
uprobs[dst_neigbours == src] = 1.0 / p
dst_src_neigbours = (np.expand_dims(dst_neigbours, 1) == np.expand_dims(
src_neigbours, 0)).sum(1).astype(np.bool8)
uprobs[dst_src_neigbours] = 1
return utils.alias_setup(uprobs / uprobs.sum())