From dc9263483017ff78fd062028ee538e12cb2c9296 Mon Sep 17 00:00:00 2001
From: Stanislaw Adaszewski <stanislaw.adaszewski@roche.com>
Date: Sun, 7 Jun 2020 15:55:59 +0200
Subject: [PATCH] Same test with torch.nn.Sequential.

---
 tests/icosagon/test_declayer.py | 50 +++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)
 create mode 100644 tests/icosagon/test_declayer.py

diff --git a/tests/icosagon/test_declayer.py b/tests/icosagon/test_declayer.py
new file mode 100644
index 0000000..fb29831
--- /dev/null
+++ b/tests/icosagon/test_declayer.py
@@ -0,0 +1,50 @@
+#
+# Copyright (C) Stanislaw Adaszewski, 2020
+# License: GPLv3
+#
+
+
+from icosagon.input import OneHotInputLayer
+from icosagon.convlayer import DecagonLayer
+from icosagon.declayer import DecodeLayer
+from icosagon.decode import DEDICOMDecoder
+from icosagon.data import Data
+import torch
+
+
+def test_decode_layer_01():
+    d = Data()
+    d.add_node_type('Dummy', 100)
+    d.add_relation_type('Dummy Relation 1', 0, 0,
+        torch.rand((100, 100), dtype=torch.float32).round().to_sparse())
+    in_layer = OneHotInputLayer(d)
+    d_layer = DecagonLayer(in_layer.output_dim, 32, d)
+    seq = torch.nn.Sequential(in_layer, d_layer)
+    last_layer_repr = seq(None)
+    dec = DecodeLayer(input_dim=d_layer.output_dim, data=d, keep_prob=1.,
+        decoder_class=DEDICOMDecoder, activation=lambda x: x)
+    pred_adj_matrices = dec(last_layer_repr)
+    assert isinstance(pred_adj_matrices, dict)
+    assert len(pred_adj_matrices) == 1
+    assert isinstance(pred_adj_matrices[0, 0], list)
+    assert len(pred_adj_matrices[0, 0]) == 1
+
+
+def test_decode_layer_02():
+    d = Data()
+    d.add_node_type('Dummy', 100)
+    d.add_relation_type('Dummy Relation 1', 0, 0,
+        torch.rand((100, 100), dtype=torch.float32).round().to_sparse())
+
+    in_layer = OneHotInputLayer(d)
+    d_layer = DecagonLayer(in_layer.output_dim, 32, d)
+    dec_layer = DecodeLayer(input_dim=d_layer.output_dim, data=d, keep_prob=1.,
+        decoder_class=DEDICOMDecoder, activation=lambda x: x)
+    seq = torch.nn.Sequential(in_layer, d_layer, dec_layer)
+
+    pred_adj_matrices = seq(None)
+    
+    assert isinstance(pred_adj_matrices, dict)
+    assert len(pred_adj_matrices) == 1
+    assert isinstance(pred_adj_matrices[0, 0], list)
+    assert len(pred_adj_matrices[0, 0]) == 1