def test_typical_node(self):
idx = 0
# edge 0 has weight 1, edge 1 has weight 2
idx2features = csr_matrix([
[1, 0], # node 0
[1, 0], # node 1
[1, 2], # node 2
[0, 2] # node 3
])
is_edge = False
actual = [
SameTypeJaccardSample((idx, 1), idx2features, is_edge),
SameTypeJaccardSample((idx, 2), idx2features, is_edge),
SameTypeJaccardSample((idx, 3), idx2features, is_edge)
]
expected = [
SimilarityRecord(left_node_idx=idx,
right_node_idx=1,
node_node_prob=1),
SimilarityRecord(left_node_idx=idx,
right_node_idx=2,
node_node_prob=1 / 3),
SimilarityRecord(left_node_idx=idx,
right_node_idx=3,
node_node_prob=0)
]
self.assertSimRecsMatch(actual, expected)
def test_node_sim_weighted(self):
_input = SimilarityRecord(left_node_idx=0,
right_node_idx=1,
left_weight=0.3,
right_weight=0.6,
node_node_prob=0.5)
actual = SamplesToModelInput([_input], num_neighbors=2, weighted=True)
# note, the node indices should be incremented
expected = (
[
[1], # left_node_idx
[0], # left_edge_idx
[2], # right_node_idx
[0], # right_edge_idx
[0.3], # left_weight
[0.6], # right_weight
[0],
[0], # neighbor_node_indices
[0],
[0], # neighbor_node_weights
[0],
[0], # neighbor_edge_indices
[0],
[0]
], # neighbor_edge_weights
[
[0.5], # node_node_prob
[0], # edge_edge_prob
[0]
]) # node_edge_prob
self.assertEqual(actual, expected)
def test_typical_edge(self):
idx = 0
idx2features = csr_matrix([
[1, 0], # edge 0
[1, 0], # edge 1
[1, 2], # edge 2
[0, 2] # edge 3
])
is_edge = True
actual = [
SameTypeJaccardSample((idx, 1), idx2features, is_edge),
SameTypeJaccardSample((idx, 2), idx2features, is_edge)
]
expected = [
SimilarityRecord(left_edge_idx=idx,
right_edge_idx=1,
edge_edge_prob=1),
SimilarityRecord(left_edge_idx=idx,
right_edge_idx=2,
edge_edge_prob=1 / 3)
]
self.assertSimRecsMatch(actual, expected)