add some tests

Author: jjcmoon

tests/test_manual.py

@@ -19,15 +19,46 @@ def test_or_node():
     assert m(weights) == 0.4 + (1 - 0.8)
 
 
-def test_probabilistic_or_node():
+def test_disjoin_conjoin():
     c = klay.Circuit()
-    l1, l2 = c.literal_node(1), c.literal_node(-2)
-    c.set_root(c.or_node([l1, l2]))
+    l1, l2, l3 = c.literal_node(1), c.literal_node(-2), c.literal_node(3)
+    or_node1 = c.disjoin([l1, l2, l2])
+    or_node2 = c.disjoin([l3, or_node1, l3, or_node1])
+    c.set_root(or_node2)
+
+    m = c.to_torch_module(semiring='real')
+    weights = torch.tensor([0.4, 0.8, 0.5])
+    expected_result = torch.tensor(0.4 + (1 - 0.8) + 0.5)
+    assert torch.allclose(m(weights), expected_result)
+
+
+def test_probabilistic():
+    c = klay.Circuit()
+    l1, l2, l3 = c.literal_node(1), c.literal_node(-2), c.literal_node(3)
+    or_node1 = c.or_node([l1, l2])
+    or_node2 = c.or_node([l2, l3])
+    and_node = c.and_node([or_node1, or_node2])
+    c.set_root(and_node)
 
     m = c.to_torch_module(semiring='real', probabilistic=True)
     m.layers[1].weights.data.zero_()
-    weights = torch.tensor([0.4, 0.8])
-    assert m(weights) == 0.5 * 0.4 + 0.5 * (1 - 0.8)
+    weights = torch.tensor([0.4, 0.8, 0.5])
+    expected_result = torch.tensor((0.4/2 + 0.2/2) * (0.2/2 + 0.5/2))
+    assert torch.allclose(m(weights), expected_result)
+
+def test_log_probabilistic():
+    c = klay.Circuit()
+    l1, l2, l3 = c.literal_node(1), c.literal_node(-2), c.literal_node(3)
+    or_node1 = c.or_node([l1, l2])
+    or_node2 = c.or_node([l2, l3])
+    and_node = c.and_node([or_node1, or_node2])
+    c.set_root(and_node)
+
+    m = c.to_torch_module(semiring='log', probabilistic=True)
+    m.layers[1].weights.data.zero_()
+    weights = torch.tensor([0.4, 0.8, 0.5])
+    expected_result = torch.tensor((0.4/2 + 0.2/2) * (0.2/2 + 0.5/2))
+    assert torch.allclose(m(weights.log()).exp(), expected_result)
 
 
 def test_multi_rooted():