From 58b0aac18b5a1c212f981543f4a3baf07d8c0b5f Mon Sep 17 00:00:00 2001
From: Stanislaw Adaszewski <stanislaw.adaszewski@roche.com>
Date: Fri, 21 Aug 2020 11:07:27 +0200
Subject: [PATCH] Fixes in sampling.

---
 src/triacontagon/sampling.py    |  4 +++-
 src/triacontagon/split.py       |  6 +++---
 tests/triacontagon/test_loop.py | 23 +++++++++++++----------
 3 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/src/triacontagon/sampling.py b/src/triacontagon/sampling.py
index 8beac99..bcec040 100644
--- a/src/triacontagon/sampling.py
+++ b/src/triacontagon/sampling.py
@@ -28,7 +28,7 @@ def fixed_unigram_candidate_sampler(
     if len(num_repeats.shape) != 1:
         raise ValueError('num_repeats must be 1D')
 
-    if torch.any(len(unigrams) - \
+    if torch.any((unigrams > 0).sum() - \
         (true_classes >= 0).sum(dim=1) < \
             num_repeats):
         raise ValueError('Not enough classes to choose from')
@@ -132,6 +132,7 @@ def negative_sample_adj_mat(adj_mat: torch.Tensor,
         raise ValueError('adj_mat must be a torch.Tensor, got: %s' % adj_mat.__class__.__name__)
 
     edges_pos, degrees = get_edges_and_degrees(adj_mat)
+    degrees = degrees.to(torch.float32) + 1.0 / torch.numel(adj_mat)
 
     true_classes, row_count = get_true_classes(adj_mat)
     if remove_diagonal:
@@ -172,6 +173,7 @@ def negative_sample_data(data: Data) -> Data:
     for vt in data.vertex_types:
         res.add_vertex_type(vt.name, vt.count)
     for key, et in data.edge_types.items():
+        print('key:', key)
         adjacency_matrices_neg = []
         for adj_mat in et.adjacency_matrices:
             remove_diagonal = True \
diff --git a/src/triacontagon/split.py b/src/triacontagon/split.py
index 77f5eaf..c105e4e 100644
--- a/src/triacontagon/split.py
+++ b/src/triacontagon/split.py
@@ -30,7 +30,7 @@ def split_adj_mat(adj_mat: torch.Tensor, ratios: List[float]):
 
         ind = indices[:, beg:end]
         val = values[beg:end]
-        res.append(_sparse_coo_tensor(ind, val, adj_mat.shape))
+        res.append(_sparse_coo_tensor(ind, val, adj_mat.shape).coalesce())
         # ofs += cnt
 
     return res
@@ -40,7 +40,7 @@ def split_edge_type(et: EdgeType, ratios: Tuple[float, float, float]):
     ratios = list(ratios)
     if sum(ratios) != 1:
         raise ValueError('Sum of ratios must be 1')
-        
+
     res = [ split_adj_mat(adj_mat, ratios) \
         for adj_mat in et.adjacency_matrices ]
 
@@ -74,7 +74,7 @@ def split_data(data: Data,
     res_1 = []
     for new_edge_types in res:
         d = Data()
-        d.vertex_types = data.vertex_types,
+        d.vertex_types = data.vertex_types
         d.edge_types = new_edge_types
         res_1.append(d)
 
diff --git a/tests/triacontagon/test_loop.py b/tests/triacontagon/test_loop.py
index dcc975a..5f4359a 100644
--- a/tests/triacontagon/test_loop.py
+++ b/tests/triacontagon/test_loop.py
@@ -78,28 +78,28 @@ def test_train_loop_01():
     data.add_vertex_type('Bar', 4)
 
     foo_foo = torch.tensor([
-        [0, 0, 0, 1, 0],
         [0, 0, 1, 0, 0],
-        [1, 0, 0, 1, 0],
-        [0, 0, 1, 0, 1],
-        [0, 1, 0, 0, 0]
+        [0, 0, 0, 1, 0],
+        [1, 0, 0, 0, 0],
+        [0, 1, 0, 0, 0],
+        [0, 0, 0, 0, 0]
     ], dtype=torch.float32)
     foo_foo = (foo_foo + foo_foo.transpose(0, 1)) / 2
 
     foo_bar = torch.tensor([
-        [0, 1, 0, 1],
+        [0, 0, 1, 0],
         [0, 0, 0, 1],
         [0, 1, 0, 0],
         [1, 0, 0, 0],
-        [0, 0, 1, 1]
+        [0, 0, 0, 1]
     ], dtype=torch.float32)
     bar_foo = foo_bar.transpose(0, 1)
 
     bar_bar = torch.tensor([
-        [0, 0, 1, 0],
-        [1, 0, 0, 0],
-        [0, 1, 0, 1],
         [0, 1, 0, 0],
+        [1, 0, 0, 0],
+        [0, 0, 0, 1],
+        [0, 0, 1, 0],
     ], dtype=torch.float32)
     bar_bar = (bar_bar + bar_bar.transpose(0, 1)) / 2
 
@@ -123,7 +123,10 @@ def test_train_loop_01():
         conv_activation=torch.sigmoid,
         dec_activation=torch.sigmoid)
 
-    train_data, val_data, test_data = split_data(data, (.9, .1, .0) )
+    train_data, val_data, test_data = split_data(data, (.5, .5, .0) )
+
+    print('val_data:', val_data)
+    print('val_data.vertex_types:', val_data.vertex_types)
 
     loop = TrainLoop(model, val_data, test_data, initial_repr,
         max_epochs=1, batch_size=1)